"""The grid — router, planner, DAG executor, and in-process sub-agents.

zeroclaw pattern: sub-agents are logical entities sharing one runtime (own workspace,
own system prompt, own tool policy) — per-agent cost is just its context window.
browser-use pattern: each browser step is an isolated Chromium agent, parallelized
with asyncio and capped by a semaphore.
"""
from __future__ import annotations

import asyncio
import json
import re
import time
from dataclasses import dataclass, field

from rich.console import Console

from .browser_worker import run_browser_task
from .cognition import (
    DENSE_TRACK,
    HUB_RULES,
    Ledger,
    MODE_GATE_RULES,
    verify_and_repair,
)
from .config import ROOT, Config
from .llm import LLM
from .memory import Memory
from .policy import Policy
from .tools import TOOL_SCHEMAS, Tools, parse_tool_args

console = Console()

# real machine facts — injected so models never hallucinate paths like C:\Users\user
import getpass
from pathlib import Path as _P
ENV_FACTS = (
    f"Machine facts: Windows username '{getpass.getuser()}', home directory {_P.home()}, "
    f"Desktop {_P.home() / 'Desktop'}. Use these EXACT paths — never invent usernames or paths."
)

AGENT_ROLES: dict[str, str] = {
    "browser": "(handled by browser-use — full Chromium autonomy)",
    "researcher": (
        "You are a fast research agent running on the user's own Windows computer with full local "
        "access. For questions about the local machine (files, apps, system state, activity), use "
        "run_shell/read_file directly. For the web, prefer http_get for static pages/APIs and "
        "browser_task for anything interactive or JS-rendered. Extract facts, cite sources."
    ),
    "coder": (
        "You are a coding agent. Write files with write_file, run/test with run_shell. "
        "Iterate until the code actually runs. Keep output minimal."
    ),
    "analyst": (
        "You are an analysis agent. Read inputs carefully, compute with run_shell (python -c) "
        "when arithmetic matters, and produce tight structured conclusions."
    ),
    "writer": (
        "You are a writing agent. Turn provided context into the requested deliverable. "
        "Use write_file if a file output is requested."
    ),
}

ROUTER_SYSTEM = (
    "You route tasks for a multi-agent system running on the user's own Windows computer with FULL "
    "local access: agents can run shell commands, read/write files anywhere, and drive a real "
    "browser. Requests about the user's machine, files, apps, or system state are normal tasks — "
    "route them to an agent (usually researcher or coder), never treat them as impossible. "
    "Reply ONLY with JSON.\n"
    'Greeting, small talk, question about the conversation itself, or anything answerable '
    'from the provided context alone → {"mode":"chat"}\n'
    'Simple single-skill task → {"mode":"direct","agent":"<browser|researcher|coder|analyst|writer>","cog":"...","task":"..."}\n'
    'Multi-part task needing decomposition → {"mode":"plan","cog":"...","task":"..."}\n'
    "Use browser for anything needing live WEB interaction; researcher for lookups; "
    "coder for writing/running code; analyst for reasoning over data; writer for prose output. "
    "IMPORTANT: the browser agent CANNOT open local files (file:// is blocked) — opening or "
    "showing a local file, folder, or locally built website routes to researcher/coder, whose "
    "show_user tool opens it on the user's screen. "
    + ENV_FACTS + " " + MODE_GATE_RULES +
    ' For direct/plan, "task" is your rewrite of the user\'s input as a fully self-contained '
    "instruction: resolve every pronoun and vague reference (it, that, the file, again) using "
    "the conversation context and memory provided — NEVER ask the user to clarify what the "
    "context can answer. Only when the referent genuinely appears nowhere in the context, "
    'route {"mode":"chat"} so the assistant can ask one precise question.'
)

CHAT_SYSTEM = (
    "You are koi-harness, a fast multi-agent assistant, talking directly with your user in a terminal. "
    "You run ON the user's Windows computer with full access to it — your agents can run any shell "
    "command, read/write any file, and drive a real browser. NEVER claim you lack access to the "
    "machine, the internet, or past activity. "
    "Answer from the conversation history and memory provided — that IS your record of what was "
    "discussed; never claim you have no memory. Resolve references yourself. Be concise and natural. "
    "If something is genuinely unresolvable from context, ask one short, precise question."
)

PLANNER_SYSTEM = (
    "You are the planner of a multi-agent grid. Decompose the task into at most {max_steps} steps. "
    "Reply ONLY with JSON: {{\"steps\":[{{\"id\":1,\"agent\":\"browser|researcher|coder|analyst|writer\","
    "\"task\":\"self-contained instruction with all context the agent needs\",\"deps\":[]}}]}}\n"
    "Steps with no unmet deps run IN PARALLEL — maximize parallelism. Each step's task must be "
    "fully self-contained (agents don't see the original prompt). Use deps only for true data "
    "dependencies; downstream steps automatically receive the outputs of their deps. "
    "Any step that combines, compares, or summarizes earlier results MUST list those step ids in its deps."
    + HUB_RULES
)


@dataclass
class Step:
    id: int
    agent: str
    task: str
    deps: list[int] = field(default_factory=list)
    result: str = ""
    done: asyncio.Event = field(default_factory=asyncio.Event)


class Orchestrator:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.llm = LLM(cfg)
        self.policy = Policy(cfg.policy_mode)
        self.memory = Memory(cfg.memory_db)
        self.history: list[tuple[str, str]] = []  # (task, answer) pairs this session
        self._opened: set = set()  # targets already presented via show_user (all runs)
        # Optional async callback(dict) the service sets to stream progress
        # (route / plan / step / step_done) to its caller as NDJSON.
        self.emit = None

    async def _e(self, kind: str, **fields) -> None:
        if self.emit is None:
            return
        try:
            await self.emit({"kind": kind, **fields})
        except Exception:
            pass  # a hung-up caller must never take the run down with it

    # extensions safe to pop open directly; code/executables open their folder instead
    _VIEWABLE = {".md", ".txt", ".html", ".htm", ".csv", ".pdf", ".png", ".jpg", ".jpeg",
                 ".gif", ".svg", ".json", ".pptx", ".xlsx", ".docx", ".mp4", ".webm", ".mp3"}

    def _open_deliverables(self, artifacts: list) -> None:
        import os
        seen: list = []
        folders: set = set()
        for p in artifacts:
            p = p.resolve()
            if str(p) in self._opened or p in seen or not p.exists():
                continue
            if "_ledger" in p.parts:
                continue
            seen.append(p)
        for p in seen[-5:]:
            try:
                if p.suffix.lower() in self._VIEWABLE:
                    os.startfile(str(p))
                    console.print(f"[dim]opened deliverable: {p.name}[/]")
                elif p.parent not in folders:
                    folders.add(p.parent)
                    os.startfile(str(p.parent))  # code/other files: show their folder
                    console.print(f"[dim]opened folder: {p.parent}[/]")
                self._opened.add(str(p))
            except OSError:
                pass

    def _context_block(self, task: str) -> str:
        parts = []
        session_tasks = {t for t, _ in self.history}
        if self.history:
            parts.append("Conversation this session (oldest first):")
            parts += [f"user: {t}\nassistant: {a[:300]}" for t, a in self.history[-6:]]
        # persistent log is ALWAYS present — "before that" must reach past session start
        older = [r for r in self.memory.recent(12)
                 if not any(r.startswith(f"earlier task: {t[:150]}") for t in session_tasks)][:8]
        if older:
            parts.append("Persistent task log from before this session (newest first):")
            parts += older
        recalls = [r for r in self.memory.recall(task, self.cfg.recall_limit) if r not in older]
        if recalls:
            parts.append("Related memory:")
            parts += recalls
        return "\n".join(parts)

    # ── entry point ──────────────────────────────────────────────
    async def handle(self, task: str) -> str:
        t0 = time.perf_counter()
        artifacts: list = []  # per-run: concurrent jobs must not open each other's files
        context = self._context_block(task)
        route = await self._route(task, context)
        mode = route.get("mode", "plan")
        cog = route.get("cog", "full")
        if cog not in ("fast", "full", "loop"):
            cog = "full"
        if self.cfg.cognition_mode in ("fast", "full", "loop"):
            cog = self.cfg.cognition_mode
        # router's self-contained rewrite: pronouns/references resolved from context
        resolved = (route.get("task") or "").strip() or task
        console.print(f"[dim]route: {mode}"
                      f"{' → ' + route.get('agent', '') if mode == 'direct' else ''}"
                      f" · cog: {cog} ({time.perf_counter() - t0:.2f}s)[/]")
        await self._e("route", mode=mode, agent=route.get("agent"), cog=cog)
        if mode == "chat":
            result = await self.llm.text(
                CHAT_SYSTEM,
                (f"Context:\n{context}\n\n" if context else "") + f"User: {task}",
                max_tokens=1024,
            )
        elif mode == "direct":
            # no verify gate here: a single agent's answer IS the only evidence, so the
            # auditor has nothing independent to check against and would demand proof
            # that cannot exist, degrading good answers into refusals
            result = await self._run_agent(route.get("agent", "researcher"), resolved, artifacts)
        else:
            result = await self._run_plan(resolved, cog, artifacts)
        self._open_deliverables(artifacts)
        if mode != "chat":  # meta-chatter would pollute recall for future sessions
            self.memory.save_run(task, result)
        self.history.append((task, result))
        console.print(f"[dim]total: {time.perf_counter() - t0:.1f}s[/]")
        return result

    async def _route(self, task: str, context: str = "") -> dict:
        user = (f"Context:\n{context}\n\n" if context else "") + f"Task: {task}"
        try:
            return await self.llm.json(ROUTER_SYSTEM, user, max_tokens=400, temperature=0.0)
        except Exception:
            return {"mode": "plan"}

    # ── planning + DAG execution ─────────────────────────────────
    async def _run_plan(self, task: str, cog: str = "full", artifacts: list | None = None) -> str:
        if artifacts is None:
            artifacts = []
        recalls = self.memory.recall(task, self.cfg.recall_limit)
        context = ("\nRelevant memory:\n" + "\n".join(recalls)) if recalls else ""
        plan = await self.llm.json(
            PLANNER_SYSTEM.format(max_steps=self.cfg.max_plan_steps),
            task + context,
            max_tokens=2048,
            temperature=0.2,
        )
        raw_steps = plan.get("steps", [])[: self.cfg.max_plan_steps]
        if not raw_steps:
            return await self._run_agent("researcher", task, artifacts)

        # broadcast hub: shared facts derived once, pushed to every step
        hub = [str(h) for h in plan.get("hub", [])][:5]
        hub_block = ("\nShared facts/constraints (derived once — do not re-derive):\n"
                     + "\n".join(f"- {h}" for h in hub)) if hub else ""
        ledger = Ledger(ROOT, task) if cog == "loop" else None
        if ledger and hub:
            ledger.set_core(hub)

        steps = {
            s["id"]: Step(
                id=s["id"],
                agent=s.get("agent", "researcher"),
                task=s.get("task", ""),
                deps=[d for d in s.get("deps", []) if d != s["id"]],
            )
            for s in raw_steps
        }
        for s in steps.values():
            s.deps = [d for d in s.deps if d in steps]
        # planner-adherence guard: a step that combines earlier results must depend on them,
        # even when the model forgot to say so
        combine_re = re.compile(r"\b(combin|compar|merge|summari|synthes|aggregat)|from step", re.IGNORECASE)
        for s in steps.values():
            earlier = [i for i in steps if i < s.id]
            if earlier and not s.deps and combine_re.search(s.task):
                s.deps = earlier
        console.print("[bold]plan:[/]")
        for s in steps.values():
            deps = f" (after {s.deps})" if s.deps else ""
            console.print(f"  [cyan]{s.id}[/] [magenta]{s.agent}[/]{deps}: {s.task[:100]}")
        await self._e(
            "plan",
            steps=[{"id": s.id, "agent": s.agent, "task": s.task[:200], "deps": s.deps} for s in steps.values()],
        )

        async def run_step(step: Step) -> None:
            for dep in step.deps:
                await steps[dep].done.wait()
            dep_context = "\n\n".join(
                f"[output of step {d}]\n{steps[d].result}" for d in step.deps if steps[d].result
            )
            full_task = step.task + hub_block + (f"\n\nContext from earlier steps:\n{dep_context}" if dep_context else "")
            console.print(f"[dim]▶ step {step.id} ({step.agent}) started[/]")
            await self._e("step", id=step.id, agent=step.agent)
            try:
                step.result = await self._run_agent(step.agent, full_task, artifacts)
            except Exception as e:
                step.result = f"step failed: {type(e).__name__}: {e}"
            console.print(f"[dim]✔ step {step.id} done[/]")
            await self._e("step_done", id=step.id, agent=step.agent, result=step.result[:300])
            step.done.set()
            if ledger:
                remaining = [s.id for s in steps.values() if not s.done.is_set()]
                ledger.seam(step.id, step.agent, step.result, remaining)

        await asyncio.gather(*(run_step(s) for s in steps.values()))

        summary = "\n\n".join(f"[step {s.id} — {s.agent}]\n{s.result}" for s in steps.values())
        ledger_block = f"\n\nRun ledger (durable state — re-read before answering):\n{ledger.render()}" if ledger else ""
        answer = await self.llm.text(
            "Synthesize the step outputs into one direct answer to the user's task. "
            "Clean, complete language only. No meta-commentary about steps or agents.",
            f"Task: {task}\n\nStep outputs:\n{summary}{ledger_block}",
            max_tokens=4096,
        )
        if cog != "fast" and self.cfg.cognition_verify:
            answer = await verify_and_repair(self.llm, task, answer, evidence=summary)
        return answer

    # ── sub-agent execution ──────────────────────────────────────
    async def _run_agent(self, agent: str, task: str, artifacts: list | None = None) -> str:
        if agent == "browser":
            return await run_browser_task(self.cfg, task)
        role = AGENT_ROLES.get(agent, AGENT_ROLES["researcher"])
        tools = Tools(self.cfg, self.policy, agent,
                      artifacts=artifacts if artifacts is not None else [], opened=self._opened)
        messages: list[dict] = [
            {"role": "system", "content": role + " " + ENV_FACTS + DENSE_TRACK +
             " If the task's deliverable is a website, URL, or finished file, call show_user "
             "with it so it appears on the user's screen — never just describe where it is."
             " Finish with a plain-text final answer."},
            {"role": "user", "content": task},
        ]
        for _ in range(self.cfg.max_worker_iterations):
            msg = await self.llm.chat(messages, tools=TOOL_SCHEMAS)
            if not msg.tool_calls:
                return msg.content or "(no answer)"
            messages.append(
                {
                    "role": "assistant",
                    "content": msg.content or "",
                    "tool_calls": [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {"name": tc.function.name, "arguments": tc.function.arguments},
                        }
                        for tc in msg.tool_calls
                    ],
                }
            )
            results = await asyncio.gather(
                *(
                    tools.dispatch(tc.function.name, parse_tool_args(tc.function.arguments))
                    for tc in msg.tool_calls
                )
            )
            for tc, res in zip(msg.tool_calls, results):
                console.print(f"[dim]  {agent} · {tc.function.name}[/]")
                messages.append({"role": "tool", "tool_call_id": tc.id, "content": res})
        return json.dumps({"warning": "iteration limit reached", "last": messages[-1]["content"][:2000]})