"""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")
        elif route == "fda":
            hits = self.client.fda(query, num_results=num_results)
        elif route == "evidence":
            hits = self.client.evidence(query, num_results=num_results)
        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 contains_cjk(text: str) -> bool:
    return any("\u4e00" <= char <= "\u9fff" for char in text)


def strip_cjk(text: str) -> str:
    return re.sub(r"[\u3400-\u9fff]+", " ", text)


def validate_packet_against_allowed_context(
    packet: dict,
    search_context: dict[str, Any] | None,
    material_context: dict[str, Any] | None,
) -> None:
    """Ensure the model did not invent source IDs or URLs beyond candidates."""
    if not search_context and not material_context:
        return
    candidates = (search_context or {}).get("candidate_sources") or []
    materials = (material_context or {}).get("materials") or []
    if not candidates and not materials:
        return
    candidate_ids = {source.get("id") for source in candidates}
    candidate_ids.update(item.get("source_id") for item in materials)
    candidate_urls = {source.get("url") for source in candidates if source.get("url")}
    candidate_urls.update(item.get("path") for item in materials if item.get("path"))
    packet_sources = packet.get("sources") or []
    unknown_ids = sorted(
        source.get("id")
        for source in packet_sources
        if source.get("id") and source.get("id") not in candidate_ids
    )
    unknown_urls = sorted(
        source.get("url")
        for source in packet_sources
        if source.get("url") and source.get("url") not in candidate_urls
    )
    if (candidates or materials) and not packet_sources:
        raise ValueError("packet must include source metadata from candidate_sources or local materials")
    if unknown_ids:
        raise ValueError(f"packet sources include non-candidate source IDs: {unknown_ids}")
    if unknown_urls:
        raise ValueError(f"packet sources include non-candidate URLs: {unknown_urls}")


def normalize_packet_against_context(
    packet: dict[str, Any],
    search_context: dict[str, Any] | None,
    material_context: dict[str, Any] | None,
) -> dict[str, Any]:
    """Deterministically fill schema metadata the model often omits."""
    packet = dict(packet)
    referenced: set[str] = set(packet.get("source_ids") or [])
    for section in ("claims", "counter_evidence"):
        for item in packet.get(section) or []:
            referenced.update(item.get("source_ids") or [])
    for item in packet.get("evidence_items") or []:
        if item.get("source_id"):
            referenced.add(item["source_id"])
    if "source_ids" not in packet or not packet.get("source_ids"):
        packet["source_ids"] = sorted(referenced)

    available_sources: dict[str, dict[str, Any]] = {}
    for source in (search_context or {}).get("candidate_sources") or []:
        if source.get("id"):
            available_sources[source["id"]] = source
    for material in (material_context or {}).get("materials") or []:
        source_id = material.get("source_id")
        if source_id:
            available_sources[source_id] = {
                "id": source_id,
                "title": material.get("title") or Path(material.get("path", "")).name,
                "url": material.get("path") or "",
                "tier": "local_material",
                "score": 8,
            }

    existing_sources = {
        source.get("id"): source
        for source in packet.get("sources") or []
        if source.get("id")
    }
    for source_id in packet.get("source_ids") or []:
        if source_id not in existing_sources and source_id in available_sources:
            existing_sources[source_id] = available_sources[source_id]
    if existing_sources:
        packet["sources"] = [existing_sources[source_id] for source_id in packet.get("source_ids", []) if source_id in existing_sources]
    return packet


FDA_AXIS_TERMS = {
    "nmpa_fda_ema_ich_who_baseline": "CGMP pharmaceutical quality system process validation aseptic processing data integrity",
    "quality_system_gap": "CGMP CAPA deviation change control data integrity quality unit pharmaceutical",
    "manufacturing_process_risk": "aseptic processing sterile drug manufacturing process validation PPQ cleaning validation water system",
    "operations_management_gap": "pharmaceutical quality system quality metrics management review senior management FDA",
    "capa_roadmap": "CGMP CAPA effectiveness remediation warning letter close-out pharmaceutical",
    "verification_evidence": "FDA 483 response CAPA effectiveness verification EIR pharmaceutical quality",
    "counter": "FDA warning letter CGMP pharmaceutical quality data integrity remediation limitations",
    "fda_enforcement_precedents": "FDA warning letter CGMP pharmaceutical aseptic processing data integrity CAPA process validation",
}


FDA_CHAPTER_TERMS = {
    "ch01": "commercial readiness phase gate remediation governance",
    "ch02": "regulatory baseline CGMP EU GMP Annex 1 ICH Q9 ICH Q10",
    "ch03": "aseptic processing RABS first air media fill visual inspection depyrogenation tunnel",
    "ch04": "biologics drug substance WFI clean utilities SCADA EMS single-use system",
    "ch05": "process validation master batch record CPP CQA PPQ cleaning validation technology transfer",
    "ch06": "deviation change control CAPA document control training data integrity quality unit",
    "ch07": "training effectiveness quality culture operator qualification human factors",
    "ch08": "quality metrics management review escalation cross-functional governance operations",
    "ch09": "CDMO quality organization technology transfer project governance capability matrix",
    "ch10": "CAPA remediation plan effectiveness check owner due date verification evidence",
    "ch11": "regulatory mapping CAPA tracker closure evidence quality assurance verification",
}


ROUTE_CHAPTER_TERMS = {
    **FDA_CHAPTER_TERMS,
}

ROUTE_SUFFIX_TERMS = {
    "scholar": "pharmaceutical GMP review validation risk management quality system",
    "patents": "biologics manufacturing patent process formulation device",
    "news": "pharmaceutical quality operations CDMO quality governance",
    "evidence": "pharmaceutical GMP evidence guidance enforcement best practice quality operations",
    "general": "pharmaceutical GMP best practice guidance quality operations remediation",
}

INTERNAL_QUERY_TOKENS = {
    "chapter_integrated",
    "input_material_findings",
}


def _compact_english_query(*parts: str, max_terms: int = 16) -> str:
    text = strip_cjk(" ".join(part for part in parts if part))
    text = re.sub(r"[^A-Za-z0-9./+-]+", " ", text)
    terms: list[str] = []
    seen: set[str] = set()
    for raw in text.split():
        term = raw.strip(" ./+-").lower()
        if not term or term in INTERNAL_QUERY_TOKENS:
            continue
        key = term.casefold()
        if key in seen:
            continue
        seen.add(key)
        terms.append(term)
        if len(terms) >= max_terms:
            break
    return " ".join(terms)


def _chapter_terms(card: TaskCard) -> str:
    mapped = " ".join(ROUTE_CHAPTER_TERMS.get(chapter_id, "") for chapter_id in card.chapter_ids)
    if mapped.strip():
        return mapped
    return strip_cjk(card.chapter_title)


def build_route_query(card: TaskCard, route: str) -> str:
    """Build short, route-aware queries instead of sending whole task cards."""
    if route == "fda":
        terms = FDA_AXIS_TERMS.get(card.topic_axis, "FDA warning letter CGMP pharmaceutical quality")
        chapter_terms = " ".join(FDA_CHAPTER_TERMS.get(chapter_id, "") for chapter_id in card.chapter_ids)
        query = f"{terms} {chapter_terms}".strip()
        if contains_cjk(query):
            raise ValueError(f"FDA route query must not contain Chinese text: {query}")
        return query
    if route == "scholar":
        return _compact_english_query(_chapter_terms(card), ROUTE_SUFFIX_TERMS["scholar"])
    if route == "patents":
        return _compact_english_query(_chapter_terms(card), ROUTE_SUFFIX_TERMS["patents"])
    if route == "news":
        return _compact_english_query(_chapter_terms(card), ROUTE_SUFFIX_TERMS["news"])
    if route == "evidence":
        return _compact_english_query(_chapter_terms(card), ROUTE_SUFFIX_TERMS["evidence"])
    return _compact_english_query(_chapter_terms(card), ROUTE_SUFFIX_TERMS["general"])


def _material_excerpt(project_root: Path | None, rel_path: str, *, max_chars: int = 6000) -> dict[str, str] | None:
    if project_root is None:
        return None
    path = project_root / rel_path
    if not path.exists() or not path.is_file():
        return None
    text = path.read_text(encoding="utf-8", errors="ignore")
    return {
        "path": rel_path,
        "source_id": f"src_local_{_safe_source_stem(Path(rel_path).stem)}",
        "title": Path(rel_path).name,
        "excerpt": text[:max_chars],
    }


def build_material_context(card: TaskCard, project_root: Path | None, *, max_chars_per_material: int = 6000) -> dict[str, Any]:
    materials = []
    seen: set[str] = set()
    for rel in card.allowed_materials:
        if rel in seen:
            continue
        seen.add(rel)
        item = _material_excerpt(project_root, rel, max_chars=max_chars_per_material)
        if item:
            materials.append(item)
    return {"materials": materials}


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
    for route in card.search_routes:
        routes_used.append(route)
        query = build_route_query(card, 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,
    material_context: dict[str, Any] | None = None,
) -> str:
    context = search_context or {"routes_used": [], "candidate_sources": []}
    materials = material_context or {"materials": []}
    return (
        "请根据以下 task card 产出一个证据包 JSON。\n"
        "正式结论、summary、open_questions 用中文；英文原文摘录、source title、DOI/URL 可以保留英文。\n"
        "必须主动包含 counter_evidence，且所有引用的 source_id 必须出现在 source_ids 中。\n\n"
        "只能使用 candidate_sources 或 Local material context 中的来源，不得编造 URL、DOI、trial ID 或 source_id。\n"
        "输出 JSON 必须包含 sources 字段；sources 只能来自 candidate_sources 或 Local material context。\n"
        "如 Local material context 非空，必须至少提取 1 条本地材料原文证据；如果与本章无关，必须在 open_questions 说明为什么无关。\n\n"
        f"{json.dumps(card.to_dict(), ensure_ascii=False, indent=2)}\n\n"
        f"Local material context:\n{json.dumps(materials, 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,
    material_context: dict[str, Any] | None = None,
) -> str:
    context = search_context or {"routes_used": [], "candidate_sources": []}
    materials = material_context or {"materials": []}
    return (
        "请修复上一次 evidence packet 输出，使其成为合法且通过 schema 校验的 JSON。\n"
        "只输出 JSON 对象，不要输出 Markdown、解释或代码块。\n"
        "保留中文主写作；英文只允许出现在来源标题、URL、DOI、原文摘录或检索笔记中。\n"
        "不得编造 candidate_sources 或 Local material context 以外的来源、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"Local material context:\n{json.dumps(materials, ensure_ascii=False, indent=2)}\n\n"
        f"Previous raw response:\n{raw_response[:12000]}"
    )


class PacketWorker:
    def __init__(
        self,
        *,
        role: RoleDefinition,
        client: ChatClient,
        project_root: Path | None = None,
        search_provider: SearchProvider | None = None,
        skill_registry: SkillRegistry | None = None,
        num_results_per_route: int = 5,
    ) -> None:
        self.role = role
        self.client = client
        self.project_root = project_root
        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 (
            f"{self.role.identity}\n\n"
            "你是 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,
            )
        material_context = build_material_context(card, self.project_root)
        raw = self.client.chat_complete(
            model=self.role.model,
            system=self._system_prompt(),
            user=build_packet_user_prompt(card, search_context, material_context),
            temperature=self.role.temperature,
            max_tokens=self.role.max_tokens,
            tag=f"packet:{card.task_id}",
        )
        try:
            packet = normalize_packet_against_context(
                _extract_json_object(raw),
                search_context,
                material_context,
            )
            validate_packet(packet)
            validate_packet_against_allowed_context(packet, search_context, material_context)
            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,
                    material_context=material_context,
                ),
                temperature=0,
                max_tokens=self.role.max_tokens,
                tag=f"packet-repair:{card.task_id}",
            )
            packet = normalize_packet_against_context(
                _extract_json_object(repaired),
                search_context,
                material_context,
            )
            validate_packet(packet)
            validate_packet_against_allowed_context(packet, search_context, material_context)
            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), project_root=project_root, 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