"""Session Context — L0 会话上下文增强（Phase 3a）

从 raw_events 重建会话状态，用于：
- 共指消解（"这个文件" → 具体路径）
- 上下文扩展词（最近编辑的文件名/实体名）

每次 recall() 调用时 from_db() 重建（~2ms），不持久化内存。
observe() hook 写入 raw_events（Issue #124），本模块只读取。
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field
from pathlib import Path

logger = logging.getLogger(__name__)


# ── 共指消解模式（Phase 3a — 不含英文 "it"）──

_COREFERENCE_PATTERNS: list[tuple[re.Pattern, str]] = [
    (re.compile(r"(这个|那个)文件"), "recent_file"),
    (re.compile(r"(这个|那个)函数"), "recent_entity"),
    (re.compile(r"(这个|那个)类"), "recent_entity"),
    (re.compile(r"(这个|那个)(?:bug|问题|错误)"), "recent_file"),
    (re.compile(r"\b(this|that) file\b", re.I), "recent_file"),
    (re.compile(r"\b(this|that) function\b", re.I), "recent_entity"),
    (re.compile(r"\b(this|that) class\b", re.I), "recent_entity"),
    (re.compile(r"\b(this|that) (?:bug|error|issue)\b", re.I), "recent_file"),
]


# ── 实体提取正则 ──

_DEF_CLASS_RE = re.compile(r"\b(?:def|class|function|fn)\s+([A-Za-z_]\w*)")

# 文件操作工具名
_FILE_TOOLS = frozenset({"Write", "Edit", "MultiEdit", "Read"})
# 代码编辑工具名
_EDIT_TOOLS = frozenset({"Edit", "Write"})


def _extract_entities_from_edit(tool_input: dict) -> list[str]:
    """从 Edit/Write 的 new_string/content 中提取 def/class 名"""
    text = tool_input.get("new_string", "")
    if not text:
        text = tool_input.get("content", "")
    if not text:
        return []
    return _DEF_CLASS_RE.findall(text)


@dataclass
class SessionContext:
    """会话上下文 — 从 raw_events 重建，不持久化内存"""

    recent_files: list[str] = field(default_factory=list)
    recent_entities: list[str] = field(default_factory=list)
    recent_queries: list[str] = field(default_factory=list)
    session_collection: str | None = None
    collection_counts: dict[str, int] = field(default_factory=dict)
    # P3-5: 本会话已返回的 fact/code IDs（降权去重用）
    returned_fact_ids: set[int] = field(default_factory=set)
    returned_code_ids: set[int] = field(default_factory=set)

    @property
    def is_warm(self) -> bool:
        """冷启动检测 — 至少 1 个文件或 2 个实体"""
        return len(self.recent_files) >= 1 or len(self.recent_entities) >= 2

    def collection_affinity(self, coll: str, primary: str) -> float:
        """返回 collection 的亲和度 [0.3, 1.0]

        - primary collection → 1.0
        - session 中访问过的 → 0.5-0.8（按访问占比线性插值）
        - 从未访问过的 → 0.3
        """
        if coll == primary:
            return 1.0
        total = sum(self.collection_counts.values()) or 1
        ratio = self.collection_counts.get(coll, 0) / total
        if ratio <= 0:
            return 0.3
        return min(0.8, 0.5 + ratio * 0.3)

    @classmethod
    def from_db(
        cls,
        db_cog: object | None,
        session_id: str | None,
        collection: str | None = None,
    ) -> SessionContext:
        """从 raw_events 重建 — ~2ms SQL 查询

        Args:
            db_cog: CogDatabase 实例（None 时返回冷启动）
            session_id: 会话 ID（None 时返回冷启动）
            collection: 限定集合（防跨项目污染）
        """
        if not session_id or not db_cog:
            return cls(session_collection=collection)
        try:
            events = db_cog.get_recent_raw_events(  # type: ignore[union-attr]
                session_id, collection=collection, limit=20,
            )
        except Exception as e:
            logger.debug("raw_events 查询失败: %s", e)
            return cls(session_collection=collection)
        # collection 访问分布（跨所有集合，~1ms）
        try:
            coll_counts = db_cog.count_session_collections(session_id)  # type: ignore[union-attr]
        except Exception:
            coll_counts = {}
        ctx = cls._parse_events(events, collection)
        ctx.collection_counts = coll_counts
        return ctx

    @classmethod
    def _parse_events(
        cls, events: list[dict], collection: str | None,
    ) -> SessionContext:
        """从 raw_events payload 提取文件和实体"""
        files: list[str] = []
        entities: list[str] = []
        seen_files: set[str] = set()
        seen_entities: set[str] = set()
        returned_fact_ids: set[int] = set()
        returned_code_ids: set[int] = set()

        for ev in events:
            raw_payload = ev.get("payload", "{}")
            try:
                payload = (
                    json.loads(raw_payload)
                    if isinstance(raw_payload, str)
                    else raw_payload
                )
            except (json.JSONDecodeError, TypeError):
                continue

            tool_name = ev.get("tool_name", "")

            # P3-5: 从 mcp_recall 事件提取已返回的 IDs
            if tool_name == "mcp_recall":
                for fid in payload.get("returned_fact_ids", []):
                    if isinstance(fid, int):
                        returned_fact_ids.add(fid)
                for cid in payload.get("returned_code_ids", []):
                    if isinstance(cid, int):
                        returned_code_ids.add(cid)
                continue  # mcp_recall 不含 tool_input

            tool_input = payload.get("tool_input", {})
            if not isinstance(tool_input, dict):
                continue

            # 文件路径提取
            if tool_name in _FILE_TOOLS:
                path = tool_input.get("file_path", "")
                if path and path not in seen_files:
                    files.append(path)
                    seen_files.add(path)

            # 实体提取（从 Edit/Write 的 new_string/content 中找 def/class）
            if tool_name in _EDIT_TOOLS:
                for ent in _extract_entities_from_edit(tool_input):
                    if ent not in seen_entities:
                        entities.append(ent)
                        seen_entities.add(ent)

        return cls(
            recent_files=files[:20],
            recent_entities=entities[:50],
            recent_queries=[],
            session_collection=collection,
            returned_fact_ids=returned_fact_ids,
            returned_code_ids=returned_code_ids,
        )

    @classmethod
    def aggregate_checkpoint(
        cls,
        db_cog: object | None,
        session_id: str | None,
    ) -> dict:
        """L0 纯规则从 raw_events 聚合会话快照 — ~5ms, 零外部依赖

        用于异常中断后的 awaken 回放：从上一个 session 的 raw_events
        重建 {active_files, recent_errors, tool_sequence} 等关键状态。

        Args:
            db_cog: CogDatabase 实例
            session_id: 要回放的 session_id
        """
        if not db_cog or not session_id:
            return {"active_files": [], "recent_entities": [],
                    "recent_errors": [], "tool_sequence": [],
                    "event_count": 0, "is_warm": False}
        try:
            events = db_cog.get_recent_raw_events(  # type: ignore[union-attr]
                session_id, limit=50,
            )
        except Exception:
            return {"active_files": [], "recent_entities": [],
                    "recent_errors": [], "tool_sequence": [],
                    "event_count": 0, "is_warm": False}

        ctx = cls._parse_events(events, None)

        recent_errors = [
            e.get("summary", "")[:100]
            for e in events
            if e.get("importance", 0) >= 0.9 and e.get("summary")
        ]

        tool_sequence = [
            e.get("tool_name", "")
            for e in reversed(events)
            if e.get("tool_name")
        ]

        return {
            "active_files": ctx.recent_files[:10],
            "recent_entities": ctx.recent_entities[:10],
            "recent_errors": recent_errors[:5],
            "tool_sequence": tool_sequence[:20],
            "event_count": len(events),
            "is_warm": ctx.is_warm,
        }

    def resolve_coreferences(self, query: str) -> str:
        """L0 共指消解 — 中文 + 英文显式指代

        只替换第一个匹配的指代词。英文 "it" 在 Phase 3b 再支持。
        """
        if not self.is_warm:
            return query

        for pattern, source_type in _COREFERENCE_PATTERNS:
            match = pattern.search(query)
            if match:
                replacement = self._get_replacement(source_type)
                if replacement:
                    query = query[:match.start()] + replacement + query[match.end():]
                    break  # 只替换第一个匹配

        return query

    def _get_replacement(self, source_type: str) -> str | None:
        """根据指代类型获取替换文本"""
        if source_type == "recent_file" and self.recent_files:
            return Path(self.recent_files[0]).name
        if source_type == "recent_entity" and self.recent_entities:
            return self.recent_entities[0]
        return None

    def get_expansion_terms(self) -> list[str]:
        """返回上下文扩展词 — 最近文件名 + 实体名"""
        if not self.is_warm:
            return []
        terms: list[str] = []
        for f in self.recent_files[:3]:
            stem = Path(f).stem
            if stem:
                terms.append(stem)
        terms.extend(self.recent_entities[:5])
        return terms
