"""Python role workers for task-card execution."""

from __future__ import annotations

import json
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Callable, Protocol, Any

from scripts.runtime.roles import RoleDefinition, RuntimeProfile
from scripts.runtime.skills import SkillRegistry
from scripts.runtime.tasks import TaskCard, validate_packet
from scripts.runtime.sources import append_packet_sources


class ChatClient(Protocol):
    def chat_complete(self, **kwargs) -> str:
        ...


class SearchProvider(Protocol):
    def search(self, *, query: str, route: str, num_results: int) -> list[dict[str, Any]]:
        ...


class ProjectSearchProvider:
    """Thin adapter over the project-owned search client."""

    def __init__(self, *, strict_specialized: bool = True) -> None:
        from scripts.lib.search_client import SearchClient

        self.client = SearchClient(strict_specialized=strict_specialized)

    def search(self, *, query: str, route: str, num_results: int) -> list[dict[str, Any]]:
        if route == "scholar":
            hits = self.client.scholar(query, num_results=num_results, year_low=2020)
        elif route == "patents":
            hits = self.client.patents(query, num_results=num_results)
        elif route == "news":
            hits = self.client.news(query, num_results=num_results, time_range="y")
        else:
            hits = self.client.search(query, num_results=num_results)
        return [
            {
                "title": hit.title,
                "url": hit.url,
                "snippet": hit.snippet,
                "route": route,
            }
            for hit in hits
        ]

    def close(self) -> None:
        self.client.close()


def _extract_json_object(text: str) -> dict:
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = stripped.strip("`")
        if stripped.startswith("json"):
            stripped = stripped[4:].strip()
    start = stripped.find("{")
    end = stripped.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("worker response does not contain a JSON object")
    return json.loads(stripped[start : end + 1])


def _safe_source_stem(task_id: str) -> str:
    return re.sub(r"[^a-zA-Z0-9]+", "_", task_id).strip("_").lower()


def build_search_context(
    card: TaskCard,
    search_provider: SearchProvider,
    *,
    num_results_per_route: int = 5,
) -> dict[str, Any]:
    candidate_sources: list[dict[str, Any]] = []
    routes_used: list[str] = []
    source_stem = _safe_source_stem(card.task_id)
    idx = 1
    query = " ".join(card.questions)
    for route in card.search_routes:
        routes_used.append(route)
        hits = search_provider.search(query=query, route=route, num_results=num_results_per_route)
        for hit in hits:
            candidate_sources.append(
                {
                    "id": f"src_{source_stem}_{idx:03d}",
                    "title": hit.get("title", ""),
                    "url": hit.get("url", ""),
                    "snippet": hit.get("snippet", ""),
                    "route": hit.get("route", route),
                    "tier": "Tier 2",
                    "score": 6,
                }
            )
            idx += 1
    return {"routes_used": routes_used, "candidate_sources": candidate_sources}


def build_packet_user_prompt(card: TaskCard, search_context: dict[str, Any] | None = None) -> str:
    context = search_context or {"routes_used": [], "candidate_sources": []}
    return (
        "请根据以下 task card 产出一个证据包 JSON。\n"
        "正式结论、summary、open_questions 用中文；英文原文摘录、source title、DOI/URL 可以保留英文。\n"
        "必须主动包含 counter_evidence，且所有引用的 source_id 必须出现在 source_ids 中。\n\n"
        "只能使用 candidate_sources 中的来源，不得编造 URL、DOI、trial ID 或 source_id。\n"
        "输出 JSON 必须包含 sources 字段，且 sources 只能来自 candidate_sources。\n\n"
        f"{json.dumps(card.to_dict(), ensure_ascii=False, indent=2)}\n\n"
        f"Search context:\n{json.dumps(context, ensure_ascii=False, indent=2)}\n\n"
        "只输出 JSON，不要输出 Markdown 解释。"
    )


def build_packet_repair_prompt(
    *,
    card: TaskCard,
    raw_response: str,
    error: Exception,
    search_context: dict[str, Any] | None = None,
) -> str:
    context = search_context or {"routes_used": [], "candidate_sources": []}
    return (
        "请修复上一次 evidence packet 输出，使其成为合法且通过 schema 校验的 JSON。\n"
        "只输出 JSON 对象，不要输出 Markdown、解释或代码块。\n"
        "保留中文主写作；英文只允许出现在来源标题、URL、DOI、原文摘录或检索笔记中。\n"
        "不得编造 candidate_sources 以外的来源、URL、DOI、trial ID 或 source_id。\n\n"
        f"Schema error:\n{error}\n\n"
        f"Task card:\n{json.dumps(card.to_dict(), ensure_ascii=False, indent=2)}\n\n"
        f"Search context:\n{json.dumps(context, ensure_ascii=False, indent=2)}\n\n"
        f"Previous raw response:\n{raw_response[:12000]}"
    )


class PacketWorker:
    def __init__(
        self,
        *,
        role: RoleDefinition,
        client: ChatClient,
        search_provider: SearchProvider | None = None,
        skill_registry: SkillRegistry | None = None,
        num_results_per_route: int = 5,
    ) -> None:
        self.role = role
        self.client = client
        self.search_provider = search_provider
        self.skill_registry = skill_registry or SkillRegistry()
        self.num_results_per_route = num_results_per_route

    def _system_prompt(self) -> str:
        skill_texts = []
        for name in self.role.skills:
            try:
                skill_texts.append(f"# Skill: {name}\n\n{self.skill_registry.read(name)}")
            except FileNotFoundError:
                skill_texts.append(f"# Skill: {name}\n\n[missing skill: {name}]")
        return (
            "你是 Deep Research v0.20 Python runtime 的证据包 worker。\n"
            "你的唯一任务是把一个 task card 转换为结构化 evidence packet。\n"
            "遵循中文主写作原则；不要写章节正文；不要编造 URL、DOI、trial ID 或 source_id。\n\n"
            "搜索只能走项目 Python search gateway 或调用方提供的 search_context；不要直接使用 Tavily MCP、browser MCP、平台 web search 或任何需要用户权限确认的外部搜索工具。\n\n"
            + "\n\n".join(skill_texts)
        )

    def run(self, card: TaskCard) -> dict:
        search_context = None
        if self.search_provider:
            search_context = build_search_context(
                card,
                self.search_provider,
                num_results_per_route=self.num_results_per_route,
            )
        raw = self.client.chat_complete(
            model=self.role.model,
            system=self._system_prompt(),
            user=build_packet_user_prompt(card, search_context),
            temperature=self.role.temperature,
            max_tokens=self.role.max_tokens,
            tag=f"packet:{card.task_id}",
        )
        try:
            packet = _extract_json_object(raw)
            validate_packet(packet)
            return packet
        except Exception as error:
            repaired = self.client.chat_complete(
                model=self.role.model,
                system=self._system_prompt(),
                user=build_packet_repair_prompt(
                    card=card,
                    raw_response=raw,
                    error=error,
                    search_context=search_context,
                ),
                temperature=0,
                max_tokens=self.role.max_tokens,
                tag=f"packet-repair:{card.task_id}",
            )
            packet = _extract_json_object(repaired)
            validate_packet(packet)
            return packet


def _write_packet_error(project_root: Path, card: TaskCard, error: Exception) -> None:
    path = project_root / "phase2" / "packet_errors" / f"{card.task_id}.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "task_id": card.task_id,
        "status": "failed",
        "error": str(error),
        "output_packet": card.output_packet,
    }
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def run_packet_workers(
    *,
    project_root: Path,
    cards: list[TaskCard],
    runtime: RuntimeProfile,
    client_factory: Callable[[RoleDefinition], ChatClient],
    search_provider_factory: Callable[[], SearchProvider] | None = None,
    workers: int,
) -> int:
    role = runtime.role_for_task("evidence_packet")
    max_workers = max(1, min(workers, role.max_concurrency))

    def run_one(card: TaskCard) -> tuple[TaskCard, dict | None, Exception | None]:
        search_provider = search_provider_factory() if search_provider_factory else None
        try:
            worker = PacketWorker(role=role, client=client_factory(role), search_provider=search_provider)
            return card, worker.run(card), None
        except Exception as error:
            return card, None, error
        finally:
            close = getattr(search_provider, "close", None)
            if close:
                close()

    written = 0
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = [pool.submit(run_one, card) for card in cards]
        for future in as_completed(futures):
            card, packet, error = future.result()
            if error is not None:
                _write_packet_error(project_root, card, error)
                continue
            if packet is None:
                _write_packet_error(project_root, card, RuntimeError("packet worker returned no packet"))
                continue
            path = project_root / card.output_packet
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(json.dumps(packet, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
            append_packet_sources(project_root / "phase2" / "sources.jsonl", packet)
            written += 1
    return written