"""Structured proposal-loop research agent — robust to small/local models, and honest about its OWN search.

Instead of free ReAct (where a weak model inconsistently decides when to call tools), this LangGraph
constrains the LLM to ONE easy job: emit a structured batch of book proposals. The graph runs each through
the portfolio gauntlet, loops, and — crucially — at the end re-judges every book it tested against the
EFFECTIVE number of independent books searched (effective_trials over their return streams). A book only
"wins" if it clears the deflated bar that accounts for the whole agentic search. The LLM proposes; the
graph + gauntlet + multiple-testing correction dispose.
"""
from __future__ import annotations

from typing import TypedDict

import numpy as np
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, Field

from fxhnt.application.portfolio_eval import PortfolioEvaluationService
from fxhnt.config import Settings
from fxhnt.domain.gauntlet import effective_trials, evaluate
from fxhnt.domain.models import AssetClass, Market, StrategySpec
from fxhnt.domain.portfolio import Book, StrategyAllocation

_SYSTEM = """You are a quantitative researcher at fxhnt. Propose diversified BOOKS to test — each a single
strategy applied across multiple markets. A deflated-Sharpe gauntlet judges them; only a PASS counts and
you cannot fake it. Prefer diversified books of 3-6 markets (the unit of alpha). Strategy-market fit:
'trend' on futures (GC NQ ES CL ZN 6E), 'buy_hold' as the baseline, 'mean_reversion' only on range-bound
markets. Use the results so far to refine toward a higher DSR. Propose 1-3 books per round."""


class BookProposal(BaseModel):
    markets: list[str] = Field(description="market symbols to combine, e.g. ['GC','NQ','ES']")
    kind: str = Field(description="strategy kind: trend | mean_reversion | buy_hold")
    window: int = Field(default=200, description="lookback window")
    rationale: str = Field(default="", description="why this book should be robust")


class ProposalBatch(BaseModel):
    reasoning: str = Field(default="", description="brief reasoning about what to test next")
    proposals: list[BookProposal] = Field(description="1-3 diversified books to test this round")


class SearchState(TypedDict):
    goal: str
    round: int
    history: list[str]
    tested: list[dict]
    proposals: list[BookProposal]


class StructuredResearchAgent:
    def __init__(self, llm, portfolio: PortfolioEvaluationService, settings: Settings, *,
                 asset_class: str = "future", max_rounds: int = 3) -> None:
        self._structured = llm.with_structured_output(ProposalBatch)
        self._portfolio = portfolio
        self._g = settings.gauntlet
        self._asset_class = AssetClass(asset_class)
        self._max_rounds = max_rounds
        self._graph = self._build()

    def _build(self):  # type: ignore[no-untyped-def]
        g = StateGraph(SearchState)
        g.add_node("propose", self._propose)
        g.add_node("evaluate", self._evaluate)
        g.add_edge(START, "propose")
        g.add_edge("propose", "evaluate")
        g.add_conditional_edges("evaluate", self._route, {"continue": "propose", "stop": END})
        return g.compile()

    def _propose(self, state: SearchState) -> dict:
        context = state["goal"] + "\n\nResults so far:\n" + ("\n".join(state["history"]) or "(none yet)")
        batch = self._structured.invoke(_SYSTEM + "\n\n" + context)
        return {"proposals": batch.proposals[:3]}

    def _evaluate(self, state: SearchState) -> dict:
        tested, history = list(state.get("tested", [])), list(state["history"])
        for p in state["proposals"]:
            try:
                markets = [Market(symbol=s.strip().upper(), asset_class=self._asset_class) for s in p.markets if s.strip()]
                if not markets or any(t["key"] == (tuple(m.symbol for m in markets), p.kind, p.window) for t in tested):
                    continue  # skip empty + already-tested (don't inflate the search with dupes)
                weight = 1.0 / len(markets)
                allocs = [StrategyAllocation(spec=StrategySpec(kind=p.kind, params={"window": float(p.window)}), market=m, weight=weight)
                          for m in markets]
                ev = self._portfolio.evaluate_book(Book(name=f"{p.kind}-book", allocations=allocs), n_book_trials=1)
                tested.append({
                    "key": (tuple(m.symbol for m in markets), p.kind, p.window),
                    "desc": f"{p.kind}(w={p.window}) on {','.join(m.symbol for m in markets)}",
                    "sharpe": ev.stats.sharpe, "oos": ev.verdict.oos_sharpe, "maxdd": ev.stats.max_drawdown,
                    "is_returns": ev.is_returns, "oos_returns": ev.oos_returns,
                })
                history.append(f"{tested[-1]['desc']}: Sharpe {ev.stats.sharpe:+.2f} OOS {ev.verdict.oos_sharpe:+.2f} (n=1 DSR {ev.verdict.dsr:.3f})")
            except Exception as e:  # surface to the loop, keep going
                history.append(f"error testing {p.markets} {p.kind}: {str(e)[:80]}")
        return {"tested": tested, "history": history, "round": state["round"] + 1}

    def _route(self, state: SearchState) -> str:
        return "stop" if state["round"] >= self._max_rounds else "continue"

    def run(self, goal: str) -> str:
        final = self._graph.invoke(
            {"goal": goal, "round": 0, "history": [], "tested": [], "proposals": []},
            {"recursion_limit": 50},
        )
        tested = final["tested"]
        if not tested:
            return "the agent proposed no testable books."

        # honest multiple-testing correction over the agent's OWN search: deflate by the EFFECTIVE number
        # of independent books it tested, not 1-per-book.
        n_eff = effective_trials([np.asarray(t["is_returns"]) for t in tested])
        n_trials = max(1, round(n_eff))
        for t in tested:
            v = evaluate(np.asarray(t["is_returns"]), np.asarray(t["oos_returns"]),
                         n_trials=n_trials, sr_variance=0.0, dsr_min=self._g.dsr_min,
                         oos_min_sharpe=self._g.oos_min_sharpe, max_is_oos_decay=self._g.max_is_oos_decay,
                         has_economic_rationale=True)
            t["dsr"], t["passed"] = v.dsr, v.passed

        tested.sort(key=lambda t: -t["dsr"])
        lines = [f"tested {len(tested)} books (≈ {n_eff:.1f} effective independent searches); "
                 f"re-judged each at the full-search bar (n_trials={n_trials}):"]
        for t in tested:
            lines.append(f"  {'PASS ' if t['passed'] else 'REJECT'} {t['desc']}: Sharpe {t['sharpe']:+.2f} "
                         f"OOS {t['oos']:+.2f} DSR {t['dsr']:.3f} maxDD {100*t['maxdd']:.0f}%")
        survivors = [t for t in tested if t["passed"]]
        best = tested[0]
        if survivors:
            lines.append(f"\nSURVIVORS ({len(survivors)}) that clear the full-search bar — best: {best['desc']} (DSR {best['dsr']:.3f}) ✅")
        else:
            lines.append(f"\nno book survives the full-search bar. Best found: {best['desc']} (DSR {best['dsr']:.3f}) — "
                         f"honest verdict: not robust once the whole search is accounted for.")
        return "\n".join(lines)
