from __future__ import annotations

import asyncio
import json
import logging
import re
import time
from dataclasses import dataclass, field
from typing import Any

from agents import (
    Agent,
    ModelSettings,
    OpenAIChatCompletionsModel,
    RunHooks,
    RunContextWrapper,
    Runner,
    function_tool,
    set_tracing_disabled,
)
from agents.exceptions import MaxTurnsExceeded
from agents.items import ModelResponse
from openai import AsyncOpenAI
from openai.types.shared import Reasoning
from pydantic import ValidationError

from .config import Settings
from .db import Database
from .schemas import CuratorDecision

logger = logging.getLogger(__name__)
PROMPT_VERSION = "maturity-2026-07-28.v3"
RECENT_CONTEXT_LIMIT = 12
IDEA_CATALOG_LIMIT = 40


CURATOR_INSTRUCTIONS = """
你是一个私人思想笔记本内部的“策展者”。用户界面必须安静，所以你的工作发生在幕后。

你的任务不是给句子贴标签，也不是强迫想法进入固定工作流，而是辨认：新片段是否属于一个
持续演化的想法；如果属于，它此刻大致在哪里、正在怎样运动、内在张力是什么。

判断“成熟度”时使用融合性的人类经验，而不是计数规则。综合考察：
1. 想法是否已经有自己的身份、边界与可复述的核心；
2. 它是否能容纳反例、摩擦、矛盾和真实经验，而不只是顺滑口号；
3. 它是否越来越属于这个具体的人，而非可替换的通用正确话；
4. 它与现实是否有接触：观察、制作、对话、试验、承诺或后果；
5. 它是否开始改变判断与行动，产生了真实影响。

0—100 只是连续的“位势坐标”，不是阶段、成绩或完成百分比。成熟度可以后退；新材料也可能
让一个看似成熟的想法重新打开。禁止用笔记数量、时间、链接数、是否列清单来机械打分。

motion 是你对当下运动的自然语言压缩，例如“向现实探去”“重构中”“暂时沉淀”，不使用
预设流水线。position 说明此刻所处的位置。trajectory 只描述最近发生的变化。possible_moves
是 1—4 个根据当前张力即时生成的可能动作，不是任务清单，不承诺固定的下一关。

谨慎合并。仅因主题词相似不能归到同一想法；关注它们是否共享同一个问题、张力或生成方向。
一个片段最多关联两个想法。若现在还不值得形成或归入想法，让它 standalone，原文仍会保留。
新建想法时 idea_id 必须为 null；更新时必须使用工具返回的准确 id。

输入会直接提供最近的连续原文和精简想法目录。把“这个、上面、第二点、继续、刚才”等指代
放回连续语境中理解；不要因为新片段换了关键词就丢掉思想上的延续。目录中的 id 只供系统内部
引用。遇到可能相关或冲突的候选时，使用 inspect_idea；search_ideas 只用于目录过大或需要
补充候选时，最多调用两次；同一个想法最多 inspect 一次。搜索无结果时不要反复改写查询，
也绝不能用 none、unknown 等占位符调用 inspect_idea。工具探索后必须及时收敛并返回最终判断。

最终必须只返回 JSON，不能用 Markdown 包裹，也不能附加解释。JSON 必须符合输入中给出的 schema。
不要在输出中评价用户、诊断心理、说教或伪造证据。
""".strip()


@dataclass
class CuratorContext:
    db: Database
    run_id: str
    user_id: str
    ideas: list[dict[str, Any]]
    recent_fragments: list[dict[str, Any]]
    search_count: int = 0
    inspection_count: int = 0
    model_rounds: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    reasoning_tokens: int = 0
    cached_tokens: int = 0
    llm_started_at: float | None = None
    inspected_idea_ids: set[str] = field(default_factory=set)


def _search_units(value: str) -> list[str]:
    units: list[str] = []
    for token in re.findall(r"[a-zA-Z0-9_-]+|[\u3400-\u9fff]+", value.lower()):
        if token not in units:
            units.append(token)
        if re.fullmatch(r"[\u3400-\u9fff]+", token) and len(token) > 2:
            for index in range(len(token) - 1):
                pair = token[index : index + 2]
                if pair not in units:
                    units.append(pair)
    return units[:24]


def _rank_ideas(
    ideas: list[dict[str, Any]], query: str
) -> tuple[list[dict[str, Any]], bool]:
    terms = _search_units(query)
    ranked: list[tuple[int, str, dict[str, Any]]] = []
    for idea in ideas:
        evidence = " ".join(
            str(fragment.get("content", ""))
            for fragment in idea.get("recent_fragments", [])
        )
        haystack = " ".join(
            str(idea.get(key, ""))
            for key in (
                "title",
                "summary",
                "position",
                "tension",
                "trajectory",
            )
        )
        haystack = f"{haystack} {evidence}".lower()
        score = sum(
            (4 if term in str(idea.get("title", "")).lower() else 1)
            * haystack.count(term)
            for term in terms
        )
        if score:
            ranked.append((score, str(idea.get("updated_at", "")), idea))
    fallback = not ranked
    if fallback:
        ranked = [
            (0, str(idea.get("updated_at", "")), idea) for idea in ideas
        ]
    ranked.sort(key=lambda item: (item[0], item[1]), reverse=True)
    return [idea for _, _, idea in ranked[:8]], fallback


class CuratorRunHooks(RunHooks[CuratorContext]):
    async def on_llm_start(
        self,
        context: RunContextWrapper[CuratorContext],
        agent: Agent[CuratorContext],
        system_prompt: str | None,
        input_items: list[Any],
    ) -> None:
        context.context.llm_started_at = time.perf_counter()

    async def on_llm_end(
        self,
        context: RunContextWrapper[CuratorContext],
        agent: Agent[CuratorContext],
        response: ModelResponse,
    ) -> None:
        state = context.context
        state.model_rounds += 1
        usage = response.usage
        input_tokens = int(usage.input_tokens or 0)
        output_tokens = int(usage.output_tokens or 0)
        reasoning_tokens = int(
            usage.output_tokens_details.reasoning_tokens or 0
        )
        cached_tokens = int(
            usage.input_tokens_details.cached_tokens or 0
        )
        state.input_tokens += input_tokens
        state.output_tokens += output_tokens
        state.reasoning_tokens += reasoning_tokens
        state.cached_tokens += cached_tokens
        duration_ms = (
            round((time.perf_counter() - state.llm_started_at) * 1000)
            if state.llm_started_at is not None
            else None
        )
        state.db.add_agent_event(
            state.run_id,
            state.user_id,
            "model_round_completed",
            {
                "round": state.model_rounds,
                "agent": agent.name,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "reasoning_tokens": reasoning_tokens,
                "cached_tokens": cached_tokens,
            },
            duration_ms=duration_ms,
        )
        state.llm_started_at = None


def build_curator_prompt(
    fragment: dict[str, Any],
    ideas: list[dict[str, Any]],
    recent_fragments: list[dict[str, Any]],
) -> str:
    prior_context = [
        {
            "id": item["id"],
            "created_at": item["created_at"],
            "content": str(item["content"])[:1_200],
        }
        for item in recent_fragments
        if item["id"] != fragment["id"]
        and item["created_at"] <= fragment["created_at"]
    ][-RECENT_CONTEXT_LIMIT:]
    later_context = [
        {
            "created_at": item["created_at"],
            "content": str(item["content"])[:1_200],
        }
        for item in recent_fragments
        if item["id"] != fragment["id"]
        and item["created_at"] > fragment["created_at"]
    ][-6:]
    catalog = [
        {
            "id": idea["id"],
            "title": idea["title"],
            "summary": idea["summary"],
            "motion": idea["motion"],
            "position": idea["position"],
            "tension": idea["tension"],
            "updated_at": idea["updated_at"],
        }
        for idea in ideas[:IDEA_CATALOG_LIMIT]
    ]
    new_fragment = {
        "id": fragment["id"],
        "created_at": fragment["created_at"],
        "content": fragment["content"],
    }
    output_schema = CuratorDecision.model_json_schema()
    catalog_note = (
        "目录已经覆盖全部已有想法。优先从目录直接判断候选，相关时用准确 id 调用 "
        "inspect_idea；一般不需要 search_ideas。"
        if len(ideas) <= IDEA_CATALOG_LIMIT
        else "目录只包含最近更新的想法；找不到可能候选时可调用 search_ideas。"
    )
    delayed_note = (
        "这是一次延迟重整。later_context 是这条记录之后才发生的输入，只用于避免把现有想法"
        "的状态和轨迹倒退；不要把它们误当成当时的上文。应把当前片段作为迟到的证据，结合"
        "想法目录中的最新状态作出增量判断。\n\n"
        if later_context
        else ""
    )
    return (
        "请策展这个刚刚保存的新片段。最近原文按时间从早到晚排列，它们是理解“这个、"
        "第二点、继续、刚才”等连续表达的第一依据；不要要求用户显式建立会话。\n\n"
        f"<recent_context>{json.dumps(prior_context, ensure_ascii=False)}</recent_context>\n\n"
        f"<new_fragment>{json.dumps(new_fragment, ensure_ascii=False)}</new_fragment>\n\n"
        f"<later_context>{json.dumps(later_context, ensure_ascii=False)}</later_context>\n\n"
        f"{delayed_note}"
        f"<idea_catalog>{json.dumps(catalog, ensure_ascii=False)}</idea_catalog>\n\n"
        f"当前共有 {len(ideas)} 个已有想法。{catalog_note}\n"
        "如果新片段既验证了一个既有想法、又形成了一个新的独立问题，可以给出两个 assessment；"
        "不要因为表面主题变化而遗漏它对上一段思想过程的反馈。\n\n"
        "最终只输出符合以下 JSON Schema 的 JSON 对象：\n"
        f"{json.dumps(output_schema, ensure_ascii=False)}"
    )


@function_tool(strict_mode=False)
async def search_ideas(
    ctx: RunContextWrapper[CuratorContext], query: str
) -> str:
    """Search existing evolving ideas by a short concept, question, or tension."""
    ctx.context.search_count += 1
    started = time.perf_counter()
    if ctx.context.search_count > 2:
        candidates = [
            {"id": item["id"], "title": item["title"]}
            for item in ctx.context.ideas[:8]
        ]
        ctx.context.db.add_agent_event(
            ctx.context.run_id,
            ctx.context.user_id,
            "tool_search_ideas",
            {
                "query": query,
                "result_count": len(candidates),
                "budget_exhausted": True,
                "candidates": candidates,
            },
            duration_ms=round((time.perf_counter() - started) * 1000),
        )
        return json.dumps(
            {
                "notice": "搜索额度已用完。使用已有目录并立即形成最终判断。",
                "candidates": candidates,
            },
            ensure_ascii=False,
        )
    ranked, fallback = _rank_ideas(ctx.context.ideas, query)
    compact = [
        {
            "id": idea["id"],
            "title": idea["title"],
            "summary": idea["summary"],
            "maturity": idea["maturity"],
            "motion": idea["motion"],
            "position": idea["position"],
            "tension": idea["tension"],
        }
        for idea in ranked
    ]
    ctx.context.db.add_agent_event(
        ctx.context.run_id,
        ctx.context.user_id,
        "tool_search_ideas",
        {
            "query": query,
            "result_count": len(compact),
            "fallback_to_recent": fallback,
            "candidates": [
                {"id": item["id"], "title": item["title"]} for item in compact
            ],
        },
        duration_ms=round((time.perf_counter() - started) * 1000),
    )
    return json.dumps(compact, ensure_ascii=False)


@function_tool(strict_mode=False)
async def inspect_idea(
    ctx: RunContextWrapper[CuratorContext], idea_id: str
) -> str:
    """Inspect an idea's recent evidence and movement before deciding to update it."""
    ctx.context.inspection_count += 1
    started = time.perf_counter()
    idea = next(
        (candidate for candidate in ctx.context.ideas if candidate["id"] == idea_id),
        None,
    )
    if not idea:
        ctx.context.db.add_agent_event(
            ctx.context.run_id,
            ctx.context.user_id,
            "tool_inspect_idea",
            {"idea_id": idea_id, "found": False},
            duration_ms=round((time.perf_counter() - started) * 1000),
        )
        return json.dumps(
            {
                "error": "idea not found",
                "valid_candidates": [
                    {"id": item["id"], "title": item["title"]}
                    for item in ctx.context.ideas[:8]
                ],
            },
            ensure_ascii=False,
        )
    if idea_id in ctx.context.inspected_idea_ids:
        ctx.context.db.add_agent_event(
            ctx.context.run_id,
            ctx.context.user_id,
            "tool_inspect_idea",
            {
                "idea_id": idea_id,
                "found": True,
                "reused": True,
                "title": idea["title"],
            },
            duration_ms=round((time.perf_counter() - started) * 1000),
        )
        return json.dumps(
            {
                "notice": (
                    "这个想法已经检查过。使用上一份结果并立即形成最终判断，"
                    "不要再次调用工具。"
                )
            },
            ensure_ascii=False,
        )
    ctx.context.inspected_idea_ids.add(idea_id)
    payload = {
        "id": idea["id"],
        "title": idea["title"],
        "summary": idea["summary"],
        "maturity_ai": idea["maturity_ai"],
        "maturity_seen_by_user": idea["maturity"],
        "user_has_overridden_position": idea["is_overridden"],
        "motion": idea["motion"],
        "position": idea["position"],
        "tension": idea["tension"],
        "trajectory": idea["trajectory"],
        "possible_moves": idea["possible_moves"],
        "recent_fragments": idea.get("recent_fragments", []),
        "recent_snapshots": idea.get("recent_snapshots", []),
    }
    ctx.context.db.add_agent_event(
        ctx.context.run_id,
        ctx.context.user_id,
        "tool_inspect_idea",
        {
            "idea_id": idea_id,
            "found": True,
            "title": idea["title"],
            "fragment_count": len(idea.get("recent_fragments", [])),
            "snapshot_count": len(idea.get("recent_snapshots", [])),
        },
        duration_ms=round((time.perf_counter() - started) * 1000),
    )
    return json.dumps(payload, ensure_ascii=False)


class Curator:
    def __init__(self, db: Database, settings: Settings):
        self.db = db
        self.settings = settings
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self._worker: asyncio.Task[None] | None = None

    @property
    def configured(self) -> bool:
        return bool(self.settings.deepseek_api_key)

    async def start(self) -> None:
        for fragment_id in self.db.pending_fragment_ids():
            self.queue.put_nowait(fragment_id)
        self._worker = asyncio.create_task(self._run_worker())

    async def stop(self) -> None:
        if self._worker:
            self._worker.cancel()
            try:
                await self._worker
            except asyncio.CancelledError:
                pass

    def enqueue(self, fragment_id: str) -> None:
        self.queue.put_nowait(fragment_id)

    async def _run_worker(self) -> None:
        while True:
            fragment_id = await self.queue.get()
            try:
                await self.analyze(fragment_id)
            except Exception as exc:
                logger.exception("Curator failed for fragment %s", fragment_id)
                self.db.set_fragment_status(
                    fragment_id, "error", str(exc)[:500]
                )
            finally:
                self.queue.task_done()

    async def analyze(self, fragment_id: str) -> None:
        fragment = self.db.get_fragment(fragment_id)
        if not fragment:
            return
        user_id = fragment.get("user_id")
        if not user_id:
            raise RuntimeError("fragment has no user owner")
        if not self.configured:
            self.db.set_fragment_status(
                fragment_id, "pending", "DeepSeek API key is not configured"
            )
            return

        self.db.set_fragment_status(fragment_id, "processing")
        run_id = self.db.start_agent_run(
            user_id,
            fragment_id,
            self.settings.deepseek_model,
            PROMPT_VERSION,
            "high",
        )
        run_started = time.perf_counter()
        ideas, recent_fragments = self.db.curator_context(user_id)
        context = CuratorContext(
            db=self.db,
            run_id=run_id,
            user_id=user_id,
            ideas=ideas,
            recent_fragments=recent_fragments,
        )
        self.db.add_agent_event(
            run_id,
            user_id,
            "context_loaded",
            {
                "fragment_characters": len(fragment["content"]),
                "idea_count": len(ideas),
                "recent_fragment_count": len(recent_fragments),
                "prior_fragment_count": sum(
                    item["id"] != fragment_id
                    and item["created_at"] <= fragment["created_at"]
                    for item in recent_fragments
                ),
                "later_fragment_count": sum(
                    item["id"] != fragment_id
                    and item["created_at"] > fragment["created_at"]
                    for item in recent_fragments
                ),
            },
        )
        prompt = build_curator_prompt(fragment, ideas, recent_fragments)

        set_tracing_disabled(True)
        client = AsyncOpenAI(
            api_key=self.settings.deepseek_api_key,
            base_url=self.settings.deepseek_base_url,
        )
        model = OpenAIChatCompletionsModel(
            model=self.settings.deepseek_model,
            openai_client=client,
        )
        model_settings = ModelSettings(
            max_tokens=2_400,
            reasoning=Reasoning(effort="high"),
            extra_body={"thinking": {"type": "enabled"}},
            extra_args={"response_format": {"type": "json_object"}},
        )
        agent: Agent[CuratorContext] = Agent(
            name="私人思想策展者",
            instructions=CURATOR_INSTRUCTIONS,
            model=model,
            tools=[search_ideas, inspect_idea],
            model_settings=model_settings,
        )
        repair_agent: Agent[CuratorContext] = Agent(
            name="判断收敛器",
            instructions=(
                CURATOR_INSTRUCTIONS
                + "\n\n这是收敛回合。不能调用任何工具；只使用输入中已经提供的连续原文、"
                "想法目录和 schema，立即给出最终合法 JSON。"
            ),
            model=model,
            tools=[],
            model_settings=model_settings,
        )
        decision: CuratorDecision | None = None
        last_error: Exception | None = None
        attempt_count = 0
        hooks = CuratorRunHooks()
        try:
            for attempt in range(2):
                attempt_count = attempt + 1
                is_repair = attempt > 0
                self.db.add_agent_event(
                    run_id,
                    user_id,
                    "attempt_started",
                    {"attempt": attempt_count, "is_repair": is_repair},
                )
                attempt_started = time.perf_counter()
                rounds_before = context.model_rounds
                usage_before = {
                    "input_tokens": context.input_tokens,
                    "output_tokens": context.output_tokens,
                    "reasoning_tokens": context.reasoning_tokens,
                    "cached_tokens": context.cached_tokens,
                }
                try:
                    result = await Runner.run(
                        repair_agent if is_repair else agent,
                        input=(
                            prompt
                            + (
                                "\n\n上一次工具探索或输出没有收敛。不要再探索，"
                                "重新完整判断并只返回合法 JSON。"
                                if is_repair
                                else ""
                            )
                        ),
                        context=context,
                        max_turns=2 if is_repair else 6,
                        hooks=hooks,
                    )
                except MaxTurnsExceeded as exc:
                    attempt_duration = round(
                        (time.perf_counter() - attempt_started) * 1000
                    )
                    current_usage = {
                        key: getattr(context, key) - value
                        for key, value in usage_before.items()
                    }
                    self.db.add_agent_event(
                        run_id,
                        user_id,
                        "model_attempt_completed",
                        {
                            "attempt": attempt_count,
                            "model_rounds": (
                                context.model_rounds - rounds_before
                            ),
                            "completed": False,
                            "error_type": type(exc).__name__,
                            **current_usage,
                        },
                        duration_ms=attempt_duration,
                    )
                    last_error = exc
                    if not is_repair:
                        self.db.add_agent_event(
                            run_id,
                            user_id,
                            "convergence_repair_started",
                            {"reason": "tool_loop_exceeded"},
                        )
                    continue
                attempt_duration = round(
                    (time.perf_counter() - attempt_started) * 1000
                )
                current_usage = {
                    key: getattr(context, key) - value
                    for key, value in usage_before.items()
                }
                self.db.add_agent_event(
                    run_id,
                    user_id,
                    "model_attempt_completed",
                    {
                        "attempt": attempt_count,
                        "model_rounds": context.model_rounds - rounds_before,
                        "completed": True,
                        **current_usage,
                    },
                    duration_ms=attempt_duration,
                )
                try:
                    raw = result.final_output
                    if not isinstance(raw, str):
                        raise TypeError("curator output was not text")
                    decision = CuratorDecision.model_validate_json(raw)
                    valid_idea_ids = {idea["id"] for idea in ideas}
                    unknown_ids = [
                        item.idea_id
                        for item in decision.assessments
                        if item.idea_id
                        and item.idea_id not in valid_idea_ids
                    ]
                    if unknown_ids:
                        raise ValueError(
                            "curator referenced an unknown existing idea"
                        )
                    break
                except (ValidationError, ValueError, TypeError) as exc:
                    decision = None
                    last_error = exc
                    self.db.add_agent_event(
                        run_id,
                        user_id,
                        "validation_failed",
                        {
                            "attempt": attempt_count,
                            "error_type": type(exc).__name__,
                            "error": str(exc)[:500],
                        },
                    )
                    logger.warning(
                        "Curator returned invalid JSON for fragment %s (attempt %s)",
                        fragment_id,
                        attempt_count,
                    )
                    if not is_repair:
                        self.db.add_agent_event(
                            run_id,
                            user_id,
                            "convergence_repair_started",
                            {"reason": "output_validation_failed"},
                        )
            if decision is None:
                raise RuntimeError(
                    "curator could not produce a valid decision"
                ) from last_error
            assessments = (
                []
                if decision.standalone
                else [item.model_dump() for item in decision.assessments]
            )
            affected = self.db.apply_assessments(
                user_id, fragment_id, assessments
            )
            self.db.add_agent_event(
                run_id,
                user_id,
                "decision_committed",
                {
                    "standalone": decision.standalone,
                    "reasoning_note": decision.reasoning_note,
                    "assessments": assessments,
                    "affected_idea_ids": affected,
                },
            )
            self.db.finish_agent_run(
                run_id,
                status="success",
                duration_ms=round((time.perf_counter() - run_started) * 1000),
                attempt_count=attempt_count,
                model_rounds=context.model_rounds,
                tool_calls=context.search_count + context.inspection_count,
                search_calls=context.search_count,
                inspection_calls=context.inspection_count,
                input_tokens=context.input_tokens,
                output_tokens=context.output_tokens,
                reasoning_tokens=context.reasoning_tokens,
                cached_tokens=context.cached_tokens,
            )
            logger.info(
                "Curator completed fragment %s with %s search and %s inspection calls",
                fragment_id,
                context.search_count,
                context.inspection_count,
            )
        except Exception as exc:
            self.db.add_agent_event(
                run_id,
                user_id,
                "run_failed",
                {
                    "error_type": type(exc).__name__,
                    "error": str(exc)[:500],
                },
            )
            self.db.finish_agent_run(
                run_id,
                status="error",
                duration_ms=round((time.perf_counter() - run_started) * 1000),
                attempt_count=attempt_count,
                model_rounds=context.model_rounds,
                tool_calls=context.search_count + context.inspection_count,
                search_calls=context.search_count,
                inspection_calls=context.inspection_count,
                error=exc,
                input_tokens=context.input_tokens,
                output_tokens=context.output_tokens,
                reasoning_tokens=context.reasoning_tokens,
                cached_tokens=context.cached_tokens,
            )
            raise