#!/usr/bin/env python
"""
CitationMemory - Global citation management system
Unified management of citation information generated by all tool calls
"""

from dataclasses import asdict, dataclass, field
from datetime import datetime
import json
from pathlib import Path
from typing import Any


@dataclass
class CitationItem:
    """Citation item"""

    cite_id: str  # Citation number, e.g., "[rag-1]"
    tool_type: str  # Tool type: rag_naive | rag_hybrid | web_search | query_item | code_execution
    query: str  # Query content
    raw_result: str = ""  # Raw tool return content (complete)
    source: str = ""  # Source information (if any)
    content: str = ""  # Citation content summary (usually summary)
    stage: str = "solve"  # Stage: analysis | solve
    step_id: str | None = None  # Belonging step (applicable to Solve stage)
    metadata: dict[str, Any] = field(default_factory=dict)  # Additional metadata
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "CitationItem":
        # Compatible with old versions without stage/step_id/metadata fields
        if "stage" not in data:
            data["stage"] = "solve"
        if "metadata" not in data:
            data["metadata"] = {}
        if "updated_at" not in data:
            data["updated_at"] = data.get("created_at", datetime.now().isoformat())
        return cls(**data)


class CitationMemory:
    """Global citation management system"""

    def __init__(self, output_dir: str | None = None):
        self.output_dir = output_dir
        self.version = "1.1"
        self.created_at = datetime.now().isoformat()
        self.updated_at = datetime.now().isoformat()

        # Core data: citation list
        self.citations: list[CitationItem] = []

        # Citation counter (incremented by tool prefix)
        self.tool_counters: dict[str, int] = {}

        # File path
        if output_dir:
            self.file_path = Path(output_dir) / "citation_memory.json"
        else:
            self.file_path = None

    @classmethod
    def load_or_create(cls, output_dir: str) -> "CitationMemory":
        """Load existing citation memory or create new memory"""
        file_path = Path(output_dir) / "citation_memory.json"

        if file_path.exists():
            # Load existing memory
            with open(file_path, encoding="utf-8") as f:
                data = json.load(f)

            memory = cls(output_dir=output_dir)
            memory.version = data.get("version", "1.1")
            memory.created_at = data.get("created_at", memory.created_at)
            memory.updated_at = data.get("updated_at", memory.updated_at)

            # Load citation list
            citations_data = data.get("citations", [])
            memory.citations = [CitationItem.from_dict(item) for item in citations_data]

            # Restore counters
            memory.tool_counters = data.get("tool_counters", {})
            if not memory.tool_counters:
                # Compatible with old versions: guess from existing cite_id
                for citation in memory.citations:
                    prefix = memory._get_tool_prefix(citation.tool_type)
                    number = memory._extract_counter_from_cite_id(citation.cite_id, prefix)
                    if number is not None:
                        memory.tool_counters[prefix] = max(
                            memory.tool_counters.get(prefix, 0), number
                        )

            return memory
        # Create new memory
        return cls(output_dir=output_dir)

    def add_citation(
        self,
        tool_type: str,
        query: str,
        raw_result: str = "",
        source: str = "",
        content: str = "",
        stage: str = "solve",
        step_id: str | None = None,
        metadata: dict[str, Any] | None = None,
        cite_id: str | None = None,
    ) -> str:
        """
        Add new citation entry

        Args:
            tool_type: Tool type
            query: Query content
            raw_result: Raw tool return content
            source: Source information
            content: Citation content summary

        Returns:
            str: Newly generated or reused cite_id (e.g., "[rag-1]")
        """
        # If cite_id is provided and already exists, return directly
        if cite_id and self.get_citation(cite_id):
            return cite_id

        # Generate new cite_id
        if not cite_id:
            cite_id = self._generate_cite_id(tool_type)

        citation = CitationItem(
            cite_id=cite_id,
            tool_type=tool_type,
            query=query,
            raw_result=raw_result,
            source=source,
            content=content,
            stage=stage,
            step_id=step_id,
            metadata=metadata or {},
        )

        self.citations.append(citation)
        self.updated_at = datetime.now().isoformat()

        return cite_id

    def get_citation(self, cite_id: str) -> CitationItem | None:
        """Get citation entry by cite_id"""
        for citation in self.citations:
            if citation.cite_id == cite_id:
                return citation
        return None

    def get_all_citations(self) -> list[CitationItem]:
        """Get all citation entries"""
        return self.citations

    def get_citations_by_tool_type(self, tool_type: str) -> list[CitationItem]:
        """Get citation entries by tool type"""
        return [c for c in self.citations if c.tool_type == tool_type]

    def update_citation(
        self,
        cite_id: str,
        raw_result: str | None = None,
        source: str | None = None,
        content: str | None = None,
        metadata: dict[str, Any] | None = None,
        stage: str | None = None,
        step_id: str | None = None,
    ):
        """Update citation entry information"""
        for citation in self.citations:
            if citation.cite_id == cite_id:
                if raw_result is not None:
                    citation.raw_result = raw_result
                if source is not None:
                    citation.source = source
                if content is not None:
                    citation.content = content
                if metadata:
                    citation.metadata.update(metadata)
                if stage:
                    citation.stage = stage
                if step_id:
                    citation.step_id = step_id
                citation.updated_at = datetime.now().isoformat()
                self.updated_at = datetime.now().isoformat()
                return
        raise ValueError(f"cite_id not found: {cite_id}")

    def save(self):
        """Save to JSON file"""
        if not self.file_path:
            raise ValueError("output_dir not set, cannot save")

        # Ensure directory exists
        self.file_path.parent.mkdir(parents=True, exist_ok=True)

        self.updated_at = datetime.now().isoformat()

        # Convert to dictionary
        data = self.to_dict()

        # Save to file
        with open(self.file_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        return {
            "version": self.version,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "citations": [item.to_dict() for item in self.citations],
            "tool_counters": self.tool_counters,
        }

    def format_citations_markdown(
        self, used_cite_ids: list[str] | None = None, language: str = "zh"
    ) -> str:
        """
        Format citation list as Markdown

        Args:
            used_cite_ids: Optional, specify citation ID list to format. If provided, only format these citations; if not provided, format all citations.
            language: Language code, 'zh' for Chinese, 'en' for English, default 'zh'

        Returns:
            str: Markdown-formatted citation list
        """
        # Set text based on language
        if language == "en":
            no_citations_text = "(No citations)"
            citation_title = "## Citations\n"
            stage_label = "Stage"
            source_label = "Source"
            query_content_label = "  - Query Content:"
            citation_content_label = "  - Citation Content:"
        else:
            no_citations_text = "(No citations)"
            citation_title = "## Citations\n"
            stage_label = "Stage"
            source_label = "Source"
            query_content_label = "  - Query Content:"
            citation_content_label = "  - Citation Content:"

        if not self.citations:
            return no_citations_text

        # If used_cite_ids is specified, only format these citations
        citations_to_format = self.citations
        if used_cite_ids is not None:
            cite_id_set = set(used_cite_ids)
            citations_to_format = [c for c in self.citations if c.cite_id in cite_id_set]
            if not citations_to_format:
                return no_citations_text

        lines = [citation_title]

        for citation in citations_to_format:
            tool_type_display = {
                "rag_naive": "RAG (Naive)",
                "rag_hybrid": "RAG (Hybrid)",
                "web_search": "Web Search",
                "query_item": "Query Item",
                "code_execution": "Code Execution",
            }.get(citation.tool_type, citation.tool_type)

            line = f"- **{citation.cite_id}** [{tool_type_display}] {stage_label}: {citation.stage}"
            if citation.source:
                line += f" | {source_label}: {citation.source}"
            lines.append(line)

            if citation.query:
                query_block = self._format_block(
                    citation.query,
                    language="python" if citation.tool_type == "code_execution" else None,
                )
                lines.append(query_content_label)
                lines.append(self._indent_block(query_block, "    "))

            if citation.content:
                content_block = self._format_block(
                    citation.content,
                    language="python" if citation.tool_type == "code_execution" else None,
                    force_block=True,
                )
                lines.append(citation_content_label)
                lines.append(self._indent_block(content_block, "    "))
            elif citation.raw_result:
                raw_block = self._format_block(
                    citation.raw_result,
                    language="python" if citation.tool_type == "code_execution" else None,
                    force_block=True,
                )
                lines.append(citation_content_label)
                lines.append(self._indent_block(raw_block, "    "))

        return "\n".join(lines)

    # ------------------------------------------------------------------
    # Internal utility methods
    # ------------------------------------------------------------------
    def _get_tool_prefix(self, tool_type: str) -> str:
        """Generate unified prefix based on tool type"""
        prefix_map = {
            "rag_naive": "rag",
            "rag_hybrid": "rag",
            "web_search": "web",
            "code_execution": "code",
            "query_item": "query",
        }
        normalized = prefix_map.get(tool_type, tool_type.replace("_", "-"))
        return normalized.lower()

    def _generate_cite_id(self, tool_type: str) -> str:
        prefix = self._get_tool_prefix(tool_type)
        current = self.tool_counters.get(prefix, 0) + 1
        self.tool_counters[prefix] = current
        return f"[{prefix}-{current}]"

    @staticmethod
    def _format_block(text: str, language: str | None = None, force_block: bool = False) -> str:
        if not text:
            return ""
        clean = text.strip()
        clean = clean.strip("`")
        if force_block or "\n" in clean or len(clean) > 80:
            lang = language or ""
            return f"```{lang}\n{clean}\n```"
        return f"`{clean}`"

    @staticmethod
    def _indent_block(block: str, prefix: str) -> str:
        if not block:
            return ""
        return "\n".join(prefix + line if line else prefix for line in block.splitlines())

    @staticmethod
    def _extract_counter_from_cite_id(cite_id: str, prefix: str) -> int | None:
        """Extract counter value from cite_id"""
        try:
            inner = cite_id.strip("[]")
            if inner.startswith(f"{prefix}-"):
                return int(inner.split("-", maxsplit=1)[1])
        except (ValueError, IndexError):
            return None
        return None
