"""Bounded program diagnostics without model input, response text or credentials."""
from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime, timezone
import hashlib
import json
import logging
import re
import uuid

from backend.orchestrator.progress_doc import task_doc_dir

log = logging.getLogger(__name__)


@dataclass
class RunDiagnostics:
    task_id: str
    run_id: str
    events: list[dict] = field(default_factory=list)
    consecutive_schema_errors: int = 0
    storage_error: bool = False


_current: ContextVar[RunDiagnostics | None] = ContextVar("browser_run_diagnostics", default=None)


@contextmanager
def run_scope(task_id: str, run_id: str):
    state = RunDiagnostics(task_id, run_id)
    token = _current.set(state)
    try:
        yield state
    finally:
        _current.reset(token)


def _schema_fields(model) -> set[str]:
    """Only program-declared field names may survive an error location."""
    fields: set[str] = set()
    try:
        schema = model.model_json_schema()
    except (AttributeError, TypeError, ValueError):
        return fields

    def visit(value):
        if isinstance(value, dict):
            fields.update(key for key in value.get("properties", {})
                          if isinstance(key, str) and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,79}", key))
            for item in value.values():
                visit(item)
        elif isinstance(value, list):
            for item in value:
                visit(item)
    visit(schema)
    return fields


def validation_fields(error, model) -> list[dict]:
    from pydantic_core._pydantic_core import list_all_errors
    known_types = {item["type"] for item in list_all_errors()}
    allowed = _schema_fields(model)
    result = []
    for item in error.errors(include_input=False, include_context=False, include_url=False)[:12]:
        # Unknown extra field names can contain a copied credential or quote.
        # Union tags and arbitrary custom error types are never logged either.
        loc = [value if isinstance(value, str) and value in allowed else value
               if type(value) is int and 0 <= value < 1000 else "<field>"
               for value in item.get("loc", ())[:12]]
        kind = item.get("type")
        result.append({"loc": loc, "type": kind if kind in known_types else "validation_error"})
    return result


def schema_error(call, error, model) -> dict:
    entry = {"kind": "schema_error", "errors": validation_fields(error, model)}
    return _record(call, entry)


def response_contract_error(call, reason: str) -> dict:
    allowed = {"missing_tool_call", "multiple_tool_calls", "unexpected_tool_call", "missing_arguments",
               "malformed_arguments", "arguments_not_object", "duplicate_argument_key", "nonfinite_arguments", "empty_actions"}
    entry = {"kind": "schema_error", "errors": [{"loc": [], "type": "response_contract"}],
             "contract": reason if reason in allowed else "invalid_response_contract"}
    return _record(call, entry)


def record_request_contract(call, function: dict) -> None:
    """Record program-owned request metadata, never the schema or messages.

    This is a prepared request, not proof of an upstream response. Join its
    call_id with model-usage response/finished events for that distinction.
    """
    if call is None or getattr(call, "phase", None) != "browser":
        return
    try:
        if (not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", call.task_id)
                or not re.fullmatch(r"[a-f0-9]{32}", call.call_id)):
            return
        definition = function.get("function", {})
        if definition.get("name") != "return_structured_output":
            return
        schema = definition["parameters"]

        def resolve(node):
            seen = set()
            while isinstance(node, dict) and "$ref" in node:
                ref = node["$ref"]
                if not isinstance(ref, str) or not ref.startswith("#/$defs/") or ref in seen:
                    return {}
                seen.add(ref)
                node = schema.get("$defs", {}).get(ref.removeprefix("#/$defs/"), {})
            return node if isinstance(node, dict) else {}

        action = resolve(schema.get("properties", {}).get("action", {}))
        items = resolve(action.get("items", {}))
        branches = items.get("anyOf", items.get("oneOf", [items]))
        names = sorted({name for branch in branches for name in resolve(branch).get("properties", {})
                        if isinstance(name, str) and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,79}", name)})
        entry = {"version": 1, "kind": "request_prepared", "task_id": call.task_id,
                 "call_id": call.call_id, "phase": call.phase,
                 "at": datetime.now(timezone.utc).isoformat(timespec="milliseconds"),
                 "function_name": definition["name"], "action_names": names,
                 "schema_sha256": hashlib.sha256(json.dumps(function, ensure_ascii=False, sort_keys=True).encode()).hexdigest()}
        directory = task_doc_dir(call.task_id)
        directory.mkdir(parents=True, exist_ok=True)
        with (directory / "model-request-contracts.jsonl").open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n")
    except (OSError, ValueError, TypeError, KeyError, AttributeError):
        # Observation failure must not change the request or add a paid retry.
        log.warning("浏览器请求工具清单未能保存，已有用量与任务资料保留")


def record_browser_action(action_name: str, status: str, *, action_id: str | None = None) -> str | None:
    """Observe the real action boundary using only names and program IDs."""
    state = _current.get()
    if state is None:
        return None
    if (status not in {"invoked", "blocked", "returned", "error"}
            or not isinstance(action_name, str) or not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,79}", action_name)):
        return None
    action_id = action_id or uuid.uuid4().hex
    if not isinstance(action_id, str) or not re.fullmatch(r"[a-f0-9]{32}", action_id):
        return None
    try:
        if (not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", state.task_id)
                or not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", state.run_id)):
            return None
        entry = {"version": 1, "kind": "browser_action", "task_id": state.task_id,
                 "run_id": state.run_id, "action_id": action_id,
                 "at": datetime.now(timezone.utc).isoformat(timespec="milliseconds"),
                 "action_name": action_name, "status": status}
        directory = task_doc_dir(state.task_id)
        directory.mkdir(parents=True, exist_ok=True)
        with (directory / "browser-actions.jsonl").open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n")
    except (OSError, ValueError, TypeError):
        state.storage_error = True
        log.warning("浏览器动作调用计数未能保存，已有用量与任务资料保留")
    return action_id


def _record(call, entry: dict) -> dict:
    state = _current.get()
    phase = getattr(call, "phase", "browser")
    if state is not None and phase == "browser":
        entry.update(run_id=state.run_id, at=datetime.now(timezone.utc).isoformat(timespec="milliseconds"))
        call_id = getattr(call, "call_id", None)
        if isinstance(call_id, str) and re.fullmatch(r"[a-f0-9]{32}", call_id):
            entry["call_id"] = call_id
        state.consecutive_schema_errors += 1
        # Bound the in-memory UI summary; the append-only private task log
        # retains every actual failed response across runs.
        state.events.append(entry)
        state.events[:] = state.events[-20:]
        try:
            if not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", state.task_id):
                raise ValueError("invalid diagnostic task")
            path = task_doc_dir(state.task_id) / "model-diagnostics.jsonl"
            with path.open("a", encoding="utf-8") as stream:
                stream.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n")
        except (OSError, ValueError):
            # A failed diagnostic write must never initiate another paid call.
            state.storage_error = True
            log.warning("浏览器结构诊断未能保存，已有用量与任务资料保留")
    return entry


def validated_response(phase: str = "browser") -> None:
    state = _current.get()
    if state is not None and phase == "browser":
        state.consecutive_schema_errors = 0


def consecutive_schema_errors() -> int:
    state = _current.get()
    return state.consecutive_schema_errors if state else 0


def pending_correction_hint() -> str:
    """Only the current run's uncorrected, sanitized error guides its next step."""
    state = _current.get()
    if not state or not state.consecutive_schema_errors or not state.events:
        return ""
    return correction_hint(state.events[-1])


def correction_hint(entry: dict) -> str:
    if entry.get("contract"):
        return "上次响应未通过结构协议校验：" + entry["contract"] + "。请仅调用return_structured_output一次，提交符合完整Schema的JSON参数；不要执行或复用上次未通过校验的动作。"
    fields = []
    for error in entry.get("errors", [])[:4]:
        location = ".".join(str(value) for value in error["loc"]) or "root"
        fields.append(f"{location} ({error['type']})")
    action_hint = ""
    if any(len(error.get("loc", [])) > 1 and error["loc"][0] == "action" and type(error["loc"][1]) is int
           for error in entry.get("errors", [])):
        action_hint = "action每个数组元素只能包含一个Schema声明的动作名键及其完整参数；多个动作分成多个数组元素。"
    return "请按现有JSON Schema纠正字段类型或必需参数：" + "；".join(fields) + "。" + action_hint + "不要执行或复用上次未通过校验的动作。"


def summary(state: RunDiagnostics) -> dict:
    return {"version": 1, "events": list(state.events),
            "consecutive_schema_errors": state.consecutive_schema_errors,
            "storage_error": state.storage_error}