from __future__ import annotations

import subprocess
import threading
from pathlib import Path
from typing import Any
from uuid import uuid4

from .dispatch import run_codex_exec
from .patch_utils import detect_patch_conflicts, extract_patch_from_stdout_jsonl, touched_files_from_patch
from .state import RunState, RunStore, RunTask
from .storage import StorageLayout, append_jsonl, safe_rmtree, utc_now_iso
from .supervisor import decide_start, prune_on_start, write_requirements_dual, write_requirements_single


def _run_git(repo_root: Path, args: list[str]) -> tuple[int, str]:
    res = subprocess.run(["git", *args], cwd=repo_root, capture_output=True, text=True)
    out = (res.stdout or "") + (res.stderr or "")
    return res.returncode, out.strip()


class Worker:
    def __init__(self, store: RunStore):
        self._store = store
        self._lock = threading.Lock()

    def start(self, *, cwd: str, query: str, options: dict[str, Any]) -> RunState:
        resolved_cwd = Path(cwd).resolve()
        decision = decide_start(cwd=resolved_cwd, query=query, options=options)
        layout = StorageLayout.for_repo(decision.repo_root)
        retention_days = int(options.get("retention_days") or 10)
        prune_on_start(repo_root=decision.repo_root, retention_days=retention_days)

        execution = options.get("execution")
        execution = execution if isinstance(execution, dict) else {}
        model = str(execution.get("model") or "gpt-5.1-codex-max")
        sandbox = str(execution.get("sandbox") or "workspace-write")
        approval_policy = str(execution.get("approval_policy") or "never")

        run_id = f"v5-{uuid4().hex}"
        paths = layout.run_paths(run_id)
        paths.run_dir.mkdir(parents=True, exist_ok=True)
        paths.artifacts_dir.mkdir(parents=True, exist_ok=True)
        paths.jobs_dir.mkdir(parents=True, exist_ok=True)

        tasks: list[RunTask] = []
        if decision.secondary_root:
            req_primary, req_secondary, contract = write_requirements_dual(
                layout=layout,
                run_id=run_id,
                query=query,
                primary_root=decision.primary_root,
                secondary_root=decision.secondary_root,
            )
            primary_slug = "primary"
            secondary_slug = "secondary"
            tasks = [
                RunTask(slug=primary_slug, project_root=str(decision.primary_root), worktree=str(layout.worktrees_dir / run_id / primary_slug)),
                RunTask(slug=secondary_slug, project_root=str(decision.secondary_root), worktree=str(layout.worktrees_dir / run_id / secondary_slug)),
            ]
            important = f"已生成需求文档 + contract；等待 coder：{req_primary.name}/{req_secondary.name}"
            _ = contract
        else:
            req = write_requirements_single(layout=layout, run_id=run_id, query=query, primary_root=decision.primary_root)
            tasks = [RunTask(slug="single", project_root=str(decision.primary_root), worktree=str(layout.worktrees_dir / run_id / "single"))]
            important = f"已生成需求文档；等待 coder：{req.name}"

        now = utc_now_iso()
        state = RunState(
            run_id=run_id,
            repo_root=str(decision.repo_root),
            cwd=str(resolved_cwd),
            query=query,
            model=model,
            sandbox=sandbox,
            approval_policy=approval_policy,
            mode=decision.mode,
            phase="agent_run",
            next_action="wait",
            wait_reason="agent_running",
            created_at=now,
            updated_at=now,
            retention_days=retention_days,
            important_event=important,
            important_event_at=now,
            tasks=tasks,
        )
        self._store.save(state)
        self._write_report(layout=layout, state=state)
        append_jsonl(
            layout.kb_dir / "run_journal.jsonl",
            {
                "ts": now,
                "run_id": run_id,
                "repo_root": str(decision.repo_root),
                "cwd": str(resolved_cwd),
                "mode": decision.mode,
                "query": query,
            },
        )

        thread = threading.Thread(target=self._run, args=(state.run_id,), name=f"squad-v5-{run_id}", daemon=True)
        thread.start()
        return state

    def _run(self, run_id: str) -> None:
        try:
            state = self._store.load(run_id)
        except Exception:
            return
        repo_root = Path(state.repo_root)
        layout = StorageLayout.for_repo(repo_root)
        paths = layout.run_paths(run_id)

        def update(**fields: Any) -> None:
            with self._lock:
                current = self._store.load(run_id)
                for k, v in fields.items():
                    setattr(current, k, v)
                current.updated_at = utc_now_iso()
                current.public_status()  # refresh brief cache
                self._store.save(current)
                self._write_report(layout=layout, state=current)

        update(important_event="已开始执行 coder", important_event_at=utc_now_iso())

        # Run tasks (single or dual) and collect patch text(s).
        patches: dict[str, str] = {}
        for task in state.tasks:
            slug = task.slug
            worktree_path = Path(task.worktree)
            worktree_path.parent.mkdir(parents=True, exist_ok=True)
            # Create worktree (detached) if not exists.
            if not worktree_path.exists():
                code, out = _run_git(repo_root, ["worktree", "add", "--detach", str(worktree_path)])
                if code != 0:
                    _ = out
                    update(phase="error", next_action="wait", wait_reason="error", important_event=f"worktree 失败: {slug}", important_event_at=utc_now_iso())
                    return

            # Best-effort sparse checkout for scope reduction: restrict to project root prefix.
            project_root = Path(task.project_root)
            try:
                rel = project_root.resolve().relative_to(repo_root.resolve())
                allow_prefix = str(rel).replace("\\", "/").strip("/")
            except Exception:
                allow_prefix = ""
            if allow_prefix:
                _run_git(worktree_path, ["sparse-checkout", "init", "--no-cone"])
                _run_git(worktree_path, ["sparse-checkout", "set", "--no-cone", "--skip-checks", f"{allow_prefix}/**"])

            requirements_hint = paths.requirements_path if slug == "single" else (paths.requirements_primary_path if slug == "primary" else paths.requirements_secondary_path)
            contract_hint = paths.contract_path if state.mode == "dual" else None
            prompt = self._build_coder_prompt(requirements_path=requirements_hint, contract_path=contract_hint)

            job_id = f"{run_id}--{slug}"
            task.status = "running"
            task.started_at = utc_now_iso()
            task.job_id = job_id
            update(tasks=state.tasks)

            res = run_codex_exec(
                job_id=job_id,
                jobs_dir=paths.jobs_dir,
                prompt=prompt,
                cwd=worktree_path,
                model=state.model,
                sandbox=state.sandbox,
                approval_policy=state.approval_policy,
                extra_config=["model_reasoning_effort=\"medium\""],
            )
            task.finished_at = utc_now_iso()
            if res.status != "completed":
                task.status = "failed"
                task.error = f"codex exit={res.exit_code}"
                update(tasks=state.tasks, phase="error", next_action="wait", wait_reason="error", important_event=f"coder 失败: {slug}", important_event_at=utc_now_iso())
                return
            patch = extract_patch_from_stdout_jsonl(res.artifacts.stdout_jsonl)
            if not patch or "diff --git " not in patch:
                task.status = "failed"
                task.error = "patch not found in stdout"
                update(tasks=state.tasks, phase="error", next_action="wait", wait_reason="error", important_event=f"未提取到 patch: {slug}", important_event_at=utc_now_iso())
                return
            task.status = "completed"
            patches[slug] = patch
            update(tasks=state.tasks, important_event=f"已拿到 patch: {slug}", important_event_at=utc_now_iso())

        # Integrate patches
        if state.mode == "dual":
            patch_a = patches.get("primary", "")
            patch_b = patches.get("secondary", "")
            overlaps = detect_patch_conflicts(patch_a, patch_b)
            if overlaps:
                update(phase="error", next_action="wait", wait_reason="error", important_event=f"patch 冲突(同文件): {len(overlaps)}", important_event_at=utc_now_iso())
                return
            final_patch = (patch_a.rstrip() + "\n\n" + patch_b.lstrip()).strip("\n") + "\n"
        else:
            final_patch = patches.get("single", "").strip("\n") + "\n"

        paths.patch_path.write_text(final_patch, encoding="utf-8")
        update(phase="done", next_action="apply", wait_reason="ready_to_apply", important_event="patch 已就绪", important_event_at=utc_now_iso())

        # Cleanup in the background; do not block readiness.
        cleanup_thread = threading.Thread(target=self._cleanup, args=(run_id,), daemon=True)
        cleanup_thread.start()

    def _cleanup(self, run_id: str) -> None:
        try:
            state = self._store.load(run_id)
        except Exception:
            return
        repo_root = Path(state.repo_root)
        layout = StorageLayout.for_repo(repo_root)
        paths = layout.run_paths(run_id)

        # Best-effort: remove worktrees and worktree dir.
        for task in state.tasks:
            try:
                _run_git(repo_root, ["worktree", "remove", "--force", str(task.worktree)])
            except Exception:
                pass
        safe_rmtree(layout.worktrees_dir / run_id)
        # Keep run artifacts; do not delete patch.

        with self._lock:
            try:
                current = self._store.load(run_id)
            except Exception:
                return
            if current.phase == "done":
                current.public_status()
                self._store.save(current)
                self._write_report(layout=layout, state=current)

    def _build_coder_prompt(self, *, requirements_path: Path, contract_path: Path | None) -> str:
        parts = [
            "You are a focused code-writing subagent.",
            "Rules:",
            "- Output MUST be a git-style patch (diff) only.",
            "- Keep changes minimal and directly related to the goal.",
            "- Self-check acceptance criteria before returning.",
            f"- Read requirements: {requirements_path}",
        ]
        if contract_path:
            parts.append(f"- If needed for cross-project sync, read/write contract: {contract_path} (<=4KB).")
        parts.append("Return a git-style patch only. No prose, no Markdown fences.")
        return "\n".join(parts) + "\n"

    def apply(self, state: RunState, *, options: dict[str, Any]) -> dict[str, Any]:
        cwd_opt = options.get("cwd")
        if isinstance(cwd_opt, str) and cwd_opt.strip():
            state = self._store.load(state.run_id, cwd=cwd_opt)
        repo_root = Path(state.repo_root)
        layout = StorageLayout.for_repo(repo_root)
        patch_path = layout.run_paths(state.run_id).patch_path
        if not patch_path.exists():
            return {"run_id": state.run_id, "applied": False, "dry_run": True, "summary": "", "error": "patch.diff not found"}
        patch_text = patch_path.read_text(encoding="utf-8", errors="replace")
        dry_run = bool(options.get("dry_run", True))
        args = ["apply", "--check"] if dry_run else ["apply"]
        proc = subprocess.run(["git", *args, str(patch_path)], cwd=repo_root, capture_output=True, text=True)
        ok = proc.returncode == 0
        summary = ""
        try:
            touched = touched_files_from_patch(patch_text)
            summary = f"files={len(touched)}"
        except Exception:
            summary = ""
        if not ok:
            err = (proc.stderr or proc.stdout or "").strip()[:2000]
            return {"run_id": state.run_id, "applied": False, "dry_run": dry_run, "summary": summary, "error": err or "git apply failed"}
        if not dry_run:
            # Record apply event in state (best-effort).
            with self._lock:
                current = self._store.load(state.run_id)
                current.important_event = "patch 已应用"
                current.important_event_at = utc_now_iso()
                current.public_status()
                self._store.save(current)
                self._write_report(layout=StorageLayout.for_repo(Path(current.repo_root)), state=current)
        return {"run_id": state.run_id, "applied": True, "dry_run": dry_run, "summary": summary}

    def _write_report(self, *, layout: StorageLayout, state: RunState) -> None:
        """
        Write a short Chinese report for users to skim.
        This is an artifact (not returned via MCP tools by default).
        """
        try:
            paths = layout.run_paths(state.run_id)
            paths.artifacts_dir.mkdir(parents=True, exist_ok=True)
            task_lines: list[str] = []
            for t in state.tasks:
                task_lines.append(f"- {t.slug}: {t.status}" + (f"（{t.error}）" if t.error else ""))
            text = "\n".join(
                [
                    "# 执行报告（简版）",
                    "",
                    f"- Run ID: {state.run_id}",
                    f"- 模式: {state.mode}",
                    f"- 阶段: {state.phase}",
                    f"- 下一步: {state.next_action}",
                    f"- 等待原因: {state.wait_reason}",
                    f"- 最近进展: {state.important_event}",
                    f"- 更新时间: {state.updated_at}",
                    "",
                    "## 子任务状态",
                    *(task_lines or ["- （无）"]),
                    "",
                    "## 产物位置",
                    f"- 需求文档（英文）: `{paths.requirements_path.name}` / `{paths.requirements_primary_path.name}` / `{paths.requirements_secondary_path.name}`",
                    f"- Patch: `{paths.patch_path.name}`",
                    "",
                ]
            )
            paths.report_path.write_text(text, encoding="utf-8")
        except Exception:
            return
