"""The agent runner: one persona researches with tools under a hard step + wall-clock
budget, then finalizes structured picks. Hand-rolled (not create_react_agent) for precise
budget control and guaranteed partial-output salvage on timeouts — the two things that
matter for unattended overnight runs.

Flow: research loop (LLM <-> tools, compaction as needed) -> finalize (structured JSON
picks). Finalize is ALWAYS attempted, even after a timeout, so a stuck agent still yields
any picks it formed rather than nothing.
"""

from __future__ import annotations

import json
import time
from datetime import datetime, timezone

from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from loguru import logger

from agents import compaction
from agents.schemas import AgentOutput, AgentStatus, Instrument, Pick
from config.personas import Persona
from config.settings import get_settings
from llm.client import build_llm
from tools import build_toolset

RESEARCH_KICKOFF = (
    "Today is {date}. This is DEEP overnight research — be exhaustive, not a quick scan. "
    "Choose {n} strong picks that fit your style.\n\n"
    "1. Call get_portfolio first (your current holdings).\n"
    "2. Build a candidate list of 8-12 tickers in your style, then investigate EACH one: "
    "get_quote, get_bars (trend/momentum), get_news (catalysts), get_fundamentals "
    "(quality/valuation), and web_search / web_fetch for context and second opinions.\n"
    "3. Use run_python for ALL math (position sizing, %s, ratios) — never do arithmetic yourself.\n"
    "4. Compare candidates, then narrow to your best {n} and pressure-test each thesis "
    "(what would make it wrong?) before finalizing.\n"
    "Work one tool call at a time. Do NOT rush to finalize — keep researching until you have "
    "genuinely vetted a broad set of names.\n"
    "You may go LONG (bet a name rises) or SHORT (bet it falls). You may trade the STOCK or a "
    "single-leg OPTION: buy a CALL for a bullish view or a PUT for a bearish view. For options, "
    "call get_option_chain(underlying, 'call' or 'put') and copy an exact contract_symbol."
)

# Appended to the kickoff when the persona has an options quota (all built-ins do).
OPTIONS_MANDATE = (
    "\n\nOPTIONS MANDATE — {opt} of your {n} picks MUST be a single-leg option trade, and it "
    "must be your HIGHEST-probability idea, not an afterthought:\n"
    "- Take your single strongest thesis and express it as a CALL (bullish) or PUT (bearish).\n"
    "- Pull get_option_chain for it and evaluate 2-3 contracts: premium, bid/ask spread, "
    "~30-45 days to expiration.\n"
    "- Use run_python to compute the breakeven price and the premium at risk. The underlying "
    "must have a specific catalyst INSIDE the contract window.\n"
    "- BUDGET CHECK: one contract costs premium x 100, and the risk system caps an option "
    "trade near 10% of your account equity (check equity via get_portfolio). If your favorite "
    "contract doesn't fit, choose a different expression that does — further strike, nearer "
    "expiry, or a cheaper underlying — your call based on the market.\n"
    "- Only submit an option pick you'd give conviction >= 0.7 after pressure-testing; "
    "low-conviction option picks are dropped by the risk system and waste a trade slot."
)

# Appended to every kickoff: how the two Python tools differ and how the sandbox behaves.
# Most historical agent errors were sandbox misuse (pandas in run_python, expecting files to
# persist, forgetting print) — this onboarding block prevents them up front.
SANDBOX_GUIDE = (
    "\n\nPYTHON TOOLS GUIDE:\n"
    "- run_python: quick exact math ONLY (stdlib: math, statistics). No pandas/yfinance here.\n"
    "- run_analysis: full scripts with pandas, numpy, yfinance, matplotlib (Agg), scipy. Each "
    "call runs in a FRESH empty directory — files do NOT persist between calls, so every "
    "script must be self-contained (re-download data each time).\n"
    "- Both: you MUST print() anything you want to see; save charts with plt.savefig('name.png')."
)

# Rotating directive nudges — cycled so a long (~2hr) research session stays PRODUCTIVE
# (does new, varied work) instead of looping "am I done?". Phrased as coverage guidance,
# not commands — the agent keeps ownership of the plan and only fills genuine gaps.
DEEPEN_NUDGES = [
    "There's time left — where is your coverage thinnest? If you've examined fewer than 8-10 "
    "candidates, screen a few MORE tickers in your style (get_bars / get_news / get_fundamentals).",
    "If you haven't yet stress-tested your favorite idea: build the BEAR case for it. What would "
    "make it wrong? Look for data that challenges the thesis (web_search + web_fetch).",
    "Consider hard numbers for your leading idea: a run_analysis script (yfinance download, "
    "momentum / moving averages / volatility) is far more convincing than a headline.",
    "Have you checked the bigger picture? Macro + sector context (rates, the sector ETF, peers) "
    "via web_search and get_bars can confirm or undercut your picks.",
    "If your top pick hasn't been compared to alternatives: line it up against 2-3 sector peers "
    "on valuation and momentum. Are you holding the best expression of the thesis?",
    "Check timing if you haven't: catalysts and dates (earnings, events via get_news / "
    "web_search), then refine entry, stop, and take-profit levels with run_python.",
]

# Takes priority over the rotation while the options mandate has had NO research yet.
OPTIONS_NUDGE = (
    "Before finalizing, remember the options mandate: your strongest thesis must also be "
    "expressed as a single-leg option. Call get_option_chain on your best idea, compare 2-3 "
    "contracts (premium, spread, ~30-45 DTE), and compute breakeven with run_python."
)

# Selected in the final stretch of the time floor: converge instead of adding breadth.
CONSOLIDATE_NUDGE = (
    "You're in the final stretch. Stop adding new names — consolidate: rank your candidates, "
    "confirm entry/stop/take numbers with run_python, and double-check the catalyst dates for "
    "the picks you intend to submit."
)

FINALIZE_PROMPT = (
    "Research complete. Return a JSON object of the form {{\"picks\": [ ... ]}} containing about "
    "{n} picks. Each pick object has these fields:\n"
    '{{"ticker": "SYM", "direction": "long", "conviction": 0.72, "thesis": "...", '
    '"catalysts": ["..."], "time_horizon": "swing", "requested_weight_pct": 0.08, '
    '"entry_type": "market", "stop_loss_pct": 0.07, "take_profit_pct": 0.15, '
    '"risk_notes": "...", "data_refs": ["..."]}}\n'
    "For an OPTION pick (long call/put) also add: \"instrument\": \"option\", "
    "\"option_right\": \"call\" or \"put\", \"contract_symbol\": \"<exact OCC symbol from "
    "get_option_chain>\", \"expiry\": \"YYYY-MM-DD\", \"strike\": <number>. Omit these for stocks.\n"
    "Rules: go long OR short as your thesis dictates (direction 'long' = bet it rises, "
    "'short' = bet it falls). requested_weight_pct is a fraction of the portfolio (0.08 = 8%). "
    "Base every number on your research. Output ONLY the JSON object."
)

# Deliberately NOT the persona prompt: the persona says "always use tools / never estimate",
# which makes a reasoning model defer to tools and emit an empty final channel at finalize
# time. A clean "research is done" system makes gpt-oss commit to the JSON answer.
FINALIZE_SYSTEM = (
    "You are finalizing stock picks from COMPLETED research. Do not ask for tools or more "
    "data — research is done. Output ONLY the requested JSON object, nothing else."
)


def _short(args: dict) -> str:
    try:
        return ", ".join(f"{k}={str(v)[:30]}" for k, v in args.items())
    except Exception:
        return str(args)[:60]


def _sanitize_tool_calls(messages: list) -> None:
    """Re-encode historical tool-call arguments to valid JSON, in place.

    gpt-oss sometimes emits a tool call (especially run_analysis with a big multi-line code
    string) whose raw `arguments` aren't valid JSON. It executes fine once, but when that
    assistant turn is echoed back in the next request, oMLX's FunctionCall validator rejects
    it with a 500 ('arguments must be valid JSON') and the whole agent dies. We rewrite the
    stored raw arguments from the cleanly-parsed args dict so every echo is valid JSON.
    """
    for m in messages:
        ak = getattr(m, "additional_kwargs", None)
        if not isinstance(ak, dict):
            continue
        raws = ak.get("tool_calls")
        if not raws:
            continue
        parsed = getattr(m, "tool_calls", None) or []
        for i, raw in enumerate(raws):
            fn = raw.get("function")
            if not isinstance(fn, dict):
                continue
            args = parsed[i]["args"] if i < len(parsed) and isinstance(parsed[i].get("args"), dict) else {}
            fn["arguments"] = json.dumps(args)


def _repair_invalid_tool_calls(ai) -> list[tuple[str, str]]:
    """Convert `invalid_tool_calls` (malformed-JSON args) into valid empty-args calls, in place.

    LangChain serializes invalid_tool_calls with their RAW malformed argument string — which
    oMLX's validator 500s on at the next request (this bypasses additional_kwargs, so the
    sanitizer above can't catch it). We convert each into a valid call with args={} and return
    (id, name) pairs so the loop can attach an error ToolMessage telling the model to retry
    with valid JSON.
    """
    import uuid

    repaired: list[tuple[str, str]] = []
    invalid = getattr(ai, "invalid_tool_calls", None) or []
    if not invalid:
        return repaired
    for itc in invalid:
        name = ((itc.get("name") or "unknown").split("<|")[0].strip()) or "unknown"
        tc_id = itc.get("id") or f"call_{uuid.uuid4().hex[:8]}"
        ai.tool_calls.append({"name": name, "args": {}, "id": tc_id, "type": "tool_call"})
        repaired.append((tc_id, name))
    ai.invalid_tool_calls = []
    return repaired


def _prompt_tokens(ai, messages: list) -> int:
    """Best available count of tokens the last prompt consumed (for compaction timing)."""
    um = getattr(ai, "usage_metadata", None) or {}
    if um.get("input_tokens"):
        return um["input_tokens"]
    tu = (getattr(ai, "response_metadata", {}) or {}).get("token_usage") or {}
    if tu.get("prompt_tokens"):
        return tu["prompt_tokens"]
    return compaction.estimate_tokens(messages)


def _extract_between(text: str, open_ch: str, close_ch: str) -> str | None:
    start, end = text.find(open_ch), text.rfind(close_ch)
    if start != -1 and end != -1 and end > start:
        return text[start : end + 1]
    return None


def parse_picks(text: str, default_stop: float, default_take: float) -> tuple[list[Pick], list[str]]:
    """Parse the finalize output into validated Picks, collecting per-item errors.

    Accepts an object {"picks": [...]}, a bare array, or those embedded in surrounding text
    / code fences. Robust to reasoning-model chatter around the JSON.
    """
    t = text.strip()
    if t.startswith("```"):
        t = t.strip("`")
        if t.lower().startswith("json"):
            t = t[4:]
        t = t.strip()

    data = None
    for candidate in (t, _extract_between(t, "{", "}"), _extract_between(t, "[", "]")):
        if not candidate:
            continue
        try:
            parsed = json.loads(candidate)
        except Exception:
            continue
        if isinstance(parsed, dict) and isinstance(parsed.get("picks"), list):
            data = parsed["picks"]
            break
        if isinstance(parsed, list):
            data = parsed
            break
    if data is None:
        return [], ["finalize: no JSON picks found in model output"]

    picks, errors = [], []
    for item in data:
        if not isinstance(item, dict):
            errors.append(f"finalize: skipped non-object item {item!r}")
            continue
        item.setdefault("stop_loss_pct", default_stop)
        item.setdefault("take_profit_pct", default_take)
        try:
            picks.append(Pick(**item))
        except Exception as e:
            errors.append(f"finalize: invalid pick {item.get('ticker', '?')}: {e}")
    return picks, errors


def _count_option_picks(picks: list[Pick]) -> int:
    return sum(1 for p in picks if p.instrument == Instrument.OPTION)


def _finalize(persona: Persona, running_summary: str, reasoning_trace: list[str]):
    """Ask the model for structured picks using a clean finalize context (not the persona's
    tool-pushing prompt), fed the research findings we gathered. Retries once on a parse
    failure, and once more if the options quota is unmet and the trace actually contains
    option-chain data to draw from. Fail-open: an unmet quota never discards the night."""
    findings = "\n".join(reasoning_trace[-50:])
    if running_summary:
        findings = f"Summary:\n{running_summary}\n\nDetails:\n{findings}"
    quota = persona.options_quota
    option_line = (
        f"\nAt least {quota} pick(s) must be OPTION picks (instrument \"option\") — your "
        "highest-conviction thesis expressed as a call or put, using an exact contract_symbol "
        "copied from your get_option_chain research. If your findings contain NO real contract "
        "symbol from get_option_chain, do NOT fabricate one — submit stock picks instead; an "
        "invented symbol is rejected and wastes the pick." if quota > 0 else ""
    )
    human = (
        f"You are {persona.display_name}. Choose picks that fit your style "
        f"(typical stop ~{persona.default_stop_pct * 100:.0f}%, "
        f"take ~{persona.default_take_pct * 100:.0f}%, horizon {persona.preferred_horizon}).\n\n"
        f"RESEARCH FINDINGS:\n{findings or '(sparse — use your best judgment)'}\n\n"
        f"{FINALIZE_PROMPT.format(n=persona.target_pick_count)}{option_line}"
    )
    llm = build_llm(max_tokens=4096)  # no tools, no json_mode (json_mode returns empty on oMLX)

    def _ask(extra: str = "") -> tuple[list[Pick], list[str]]:
        resp = llm.invoke(
            [SystemMessage(content=FINALIZE_SYSTEM), HumanMessage(content=human + extra)]
        )
        content = resp.content if isinstance(resp.content, str) else str(resp.content)
        return parse_picks(content, persona.default_stop_pct, persona.default_take_pct)

    errors: list[str] = []
    picks: list[Pick] = []
    for _ in range(2):
        try:
            picks, errs = _ask()
            if picks:
                errors = errs
                break
            errors = errs
        except Exception as e:
            errors = [f"finalize llm error: {e}"]
    if not picks:
        return [], errors

    # options quota: one corrective retry, only if the trace has real contracts to use
    if quota > 0 and _count_option_picks(picks) < quota:
        trace_has_chain = any("get_option_chain" in line for line in reasoning_trace)
        if trace_has_chain:
            try:
                retry_picks, retry_errs = _ask(
                    f"\n\nYour previous answer had too few option picks ({_count_option_picks(picks)}"
                    f"/{quota} required). Re-issue ALL picks, converting your strongest thesis into "
                    "a single-leg option pick using an exact contract_symbol that appears in the "
                    "research findings. Do not invent contract symbols."
                )
                if retry_picks and _count_option_picks(retry_picks) > _count_option_picks(picks):
                    return retry_picks, retry_errs
            except Exception as e:
                errors.append(f"options-quota retry failed: {e}")
        errors.append(
            f"options mandate unmet: {_count_option_picks(picks)}/{quota} option picks (fail-open)"
        )
    return picks, errors


def run_agent(
    persona: Persona,
    *,
    deadline_ts: float,
    step_budget: int,
    min_steps: int = 0,
    min_seconds: float = 0.0,
    context_blurb: str = "",
    checkpointer=None,
    model_name: str = "",
) -> AgentOutput:
    start = time.time()
    started_at = datetime.now(timezone.utc)
    tools = build_toolset(persona)
    tools_by_name = {t.name: t for t in tools}
    tool_llm = build_llm(tools=tools)
    plain_llm = build_llm()  # no tools — used for summarization
    # Compaction ceiling is the MEMORY-safe KV budget, not the model's 131k window: gpt-oss
    # hits oMLX's prefill memory guard around ~47k tokens on a 24GB Mac, so we compact well
    # below that to keep long (2hr) runs alive.
    compact_budget = get_settings().compaction_token_budget

    date_str = datetime.now().astimezone().strftime("%A %Y-%m-%d")
    kickoff = RESEARCH_KICKOFF.format(date=date_str, n=persona.target_pick_count)
    if persona.options_quota > 0:
        kickoff += OPTIONS_MANDATE.format(opt=persona.options_quota, n=persona.target_pick_count)
    kickoff += SANDBOX_GUIDE
    if context_blurb:  # market regime + the agent's own track record (self-learning)
        kickoff = f"{context_blurb}\n\n{kickoff}"
    messages: list = [
        SystemMessage(content=persona.system_prompt),
        HumanMessage(content=kickoff),
    ]
    running_summary = ""
    reasoning_trace: list[str] = []
    steps_used = 0
    nudge_count = 0
    ended_reason = "natural"

    def snapshot(status: AgentStatus, picks: list[Pick] | None = None, errors=None) -> AgentOutput:
        return AgentOutput(
            agent_name=persona.display_name,
            persona_id=persona.id,
            started_at=started_at,
            ended_at=datetime.now(timezone.utc),
            status=status,
            steps_used=steps_used,
            wall_seconds=round(time.time() - start, 1),
            market_summary=running_summary,
            picks=picks or [],
            reasoning_trace=list(reasoning_trace),
            errors=errors or [],
        )

    # --- research loop ---
    while steps_used < step_budget:
        if time.time() >= deadline_ts:
            ended_reason = "timeout"
            break
        steps_used += 1
        _sanitize_tool_calls(messages)  # keep echoed tool-call args valid JSON (oMLX 500 guard)
        ai = None
        for attempt in range(4):
            try:
                ai = tool_llm.invoke(messages)
                break
            except Exception as e:
                err = str(e)[:100]
                reasoning_trace.append(f"[retry {attempt + 1}/4] invoke failed: {err}")
                logger.warning("{}: invoke failed (attempt {}/4): {}", persona.id, attempt + 1, err)
                time.sleep(4)
                if attempt >= 1:  # watchdog: restart oMLX if it's wedged, then keep going
                    from llm import server as llm_server

                    if not llm_server.server_up():
                        logger.warning("{}: oMLX down — watchdog restarting it", persona.id)
                        llm_server.ensure_ready(warm=False)
        if ai is None:
            reasoning_trace.append("[error] llm.invoke failed after 4 retries")
            ended_reason = "error"
            logger.warning("{}: giving up research after 4 failed invokes (will still finalize)", persona.id)
            break
        messages.append(ai)

        # malformed-JSON tool calls: make them serializable + tell the model to retry
        repaired = _repair_invalid_tool_calls(ai)
        repaired_ids = {tc_id for tc_id, _ in repaired}
        for tc_id, name in repaired:
            messages.append(
                ToolMessage(
                    content=f"ERROR: your {name} call had malformed JSON arguments. "
                    "Call it again — arguments must be a valid JSON object.",
                    tool_call_id=tc_id,
                )
            )
            reasoning_trace.append(f"[repair] malformed {name} args -> asked model to retry")

        tool_calls = getattr(ai, "tool_calls", None)
        if not tool_calls:
            # Force depth: if the model tries to finalize before the TIME floor (or step floor),
            # inject a directive nudge so it keeps doing NEW research. Bounded by deadline.
            # Priority: unmet options mandate first, converge in the final stretch, else rotate.
            elapsed = time.time() - start
            if (elapsed < min_seconds or steps_used < min_steps) and time.time() < deadline_ts:
                options_pending = persona.options_quota > 0 and not any(
                    "get_option_chain" in line for line in reasoning_trace
                )
                if options_pending:
                    nudge = OPTIONS_NUDGE
                elif min_seconds and elapsed > 0.8 * min_seconds:
                    nudge = CONSOLIDATE_NUDGE
                else:
                    nudge = DEEPEN_NUDGES[nudge_count % len(DEEPEN_NUDGES)]
                nudge_count += 1
                messages.append(HumanMessage(content=nudge))
                reasoning_trace.append(
                    f"[nudge #{nudge_count}] {elapsed/60:.0f}/{min_seconds/60:.0f}min elapsed, step {steps_used}"
                )
                continue
            if ai.content:
                reasoning_trace.append(f"[note] {str(ai.content)[:200]}")
            ended_reason = "natural"
            break

        for tc in tool_calls:
            if tc["id"] in repaired_ids:  # already answered with a retry-error ToolMessage
                continue
            # gpt-oss can leak harmony channel tokens into the tool name (e.g. "get_bars<|channel|>")
            name = tc["name"].split("<|")[0].strip()
            args = tc.get("args", {}) or {}
            tool = tools_by_name.get(name)
            try:
                result = tool.invoke(args) if tool else f"ERROR: unknown tool {name}"
            except Exception as e:
                result = f"ERROR: {type(e).__name__}: {e}"
            messages.append(ToolMessage(content=str(result), tool_call_id=tc["id"]))
            reasoning_trace.append(f"{name}({_short(args)}) -> {str(result)[:160]}")

        if compaction.should_compact(_prompt_tokens(ai, messages), compact_budget):
            messages, running_summary = compaction.compact(plain_llm, messages, running_summary)

        if checkpointer:  # incremental save after each round
            checkpointer.write(snapshot(AgentStatus.PARTIAL))

    if steps_used >= step_budget and ended_reason == "natural":
        ended_reason = "budget"

    # --- finalize (ALWAYS attempted — even after invoke failures, salvage the findings) ---
    picks: list[Pick] = []
    errors: list[str] = []
    try:
        picks, errors = _finalize(persona, running_summary, reasoning_trace)
    except Exception as e:
        errors = [f"finalize crashed: {e}"]

    if ended_reason == "error":
        # salvaged picks still count — mark partial so execution can use them
        status = AgentStatus.PARTIAL if picks else AgentStatus.ERROR
    elif ended_reason == "timeout":
        status = AgentStatus.TIMEOUT
    elif picks:
        status = AgentStatus.COMPLETE
    else:
        status = AgentStatus.PARTIAL

    out = snapshot(status, picks=picks, errors=errors)
    if checkpointer:
        checkpointer.write(out)
    logger.info(
        "{}: status={} steps={} picks={} wall={}s",
        persona.id, status.value, steps_used, len(picks), out.wall_seconds,
    )
    return out