#!/usr/bin/env python
"""
InvestigateMemory - Analysis loop memory file system (Refactored: uses unified cite_id)
"""

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


@dataclass
class KnowledgeItem:
    """Knowledge item (Refactored: uses cite_id)"""

    cite_id: str  # Citation ID, corresponds to cite_id in CitationMemory
    tool_type: str  # rag_naive | rag_hybrid | web_search | query_item
    query: str
    raw_result: str  # Raw tool return content (complete)
    summary: str = ""  # Summary generated by NoteAgent
    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]) -> "KnowledgeItem":
        # Backward compatibility: handle old version field names
        if "knowledge_id" in data and "cite_id" not in data:
            # Convert old knowledge_id to cite_id
            data["cite_id"] = data.pop("knowledge_id")
        if "source_type" in data and "tool_type" not in data:
            data["tool_type"] = data.pop("source_type")
        if "query_text" in data and "query" not in data:
            data["query"] = data.pop("query_text")
        if "answer_raw" in data and "raw_result" not in data:
            data["raw_result"] = data.pop("answer_raw")
        # Remove old citations field (now managed by CitationMemory)
        if "citations" in data:
            data.pop("citations")
        return cls(**data)


@dataclass
class Reflections:
    """Reflections"""

    remaining_questions: list[str] = field(
        default_factory=list
    )  # Questions still needing investigation
    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]) -> "Reflections":
        return cls(**data)


class InvestigateMemory:
    """Analysis loop memory management (Refactored: uses unified cite_id)"""

    def __init__(
        self,
        task_id: str | None = None,
        user_question: str = "",
        output_dir: str | None = None,
    ):
        self.task_id = task_id or f"investigate_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.user_question = user_question
        self.output_dir = output_dir

        self.version = "3.0"  # New version
        self.created_at = datetime.now().isoformat()
        self.updated_at = datetime.now().isoformat()

        # Core data
        self.knowledge_chain: list[KnowledgeItem] = []
        self.reflections: Reflections = Reflections()

        # Metadata (for statistics and monitoring)
        self.metadata: dict[str, Any] = {
            "total_iterations": 0,
            "coverage_rate": 0.0,
            "avg_confidence": 0.0,
            "total_knowledge_items": 0,
        }

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

    @classmethod
    def load_or_create(
        cls, output_dir: str, user_question: str = "", task_id: str | None = None
    ) -> "InvestigateMemory":
        """Load existing memory or create new memory (supports v1.0/v2.0 backward compatibility)"""
        file_path = Path(output_dir) / "investigate_memory.json"

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

            memory = cls(
                task_id=data.get("task_id"),
                user_question=data.get("user_question", user_question),
                output_dir=output_dir,
            )

            file_version = data.get("version", "1.0")
            memory.version = "3.0"  # Upgrade to new version
            memory.created_at = data.get("created_at", memory.created_at)
            memory.updated_at = data.get("updated_at", memory.updated_at)

            # Load knowledge chain (supports v1.0/v2.0 compatibility)
            knowledge_chain_data = data.get("knowledge_chain", [])
            memory.knowledge_chain = [
                KnowledgeItem.from_dict(item) for item in knowledge_chain_data
            ]

            # If v1.0, need to migrate data
            if file_version == "1.0":
                # Merge notes summary into knowledge_chain
                notes_data = data.get("notes", [])
                for note in notes_data:
                    related_knowledge_ids = note.get("related_knowledge_ids", [])

                    # Prioritize using related_knowledge_ids
                    for knowledge_id in related_knowledge_ids:
                        for k_item in memory.knowledge_chain:
                            # Old version uses knowledge_id, new version uses cite_id
                            if k_item.cite_id == knowledge_id:
                                if not k_item.summary:
                                    k_item.summary = note.get("summary", "")
                                break

                # Convert reflections to remaining_questions
                reflections_data = data.get("reflections", [])
                remaining_questions = []
                for reflection in reflections_data:
                    action_items = reflection.get("action_items", [])
                    remaining_questions.extend(action_items)
                    gaps = reflection.get("gaps", [])
                    remaining_questions.extend(gaps)

                # Deduplicate
                memory.reflections.remaining_questions = list(set(remaining_questions))
            else:
                # v2.0/v3.0 format
                reflections_data = data.get("reflections", {})
                if isinstance(reflections_data, dict):
                    memory.reflections = Reflections.from_dict(reflections_data)
                else:
                    memory.reflections = Reflections()

            # Load metadata (if exists)
            memory.metadata = data.get("metadata", memory.metadata)

            return memory
        # Create new memory
        return cls(task_id=task_id, user_question=user_question, output_dir=output_dir)

    def add_knowledge(self, item: KnowledgeItem):
        """Add knowledge item"""
        self.knowledge_chain.append(item)
        self.updated_at = datetime.now().isoformat()

    def update_knowledge_summary(self, cite_id: str, summary: str):
        """Update knowledge item summary (called by NoteAgent)"""
        for item in self.knowledge_chain:
            if item.cite_id == cite_id:
                item.summary = summary
                item.updated_at = datetime.now().isoformat()
                self.updated_at = datetime.now().isoformat()
                return
        raise ValueError(f"cite_id not found: {cite_id}")

    def get_available_knowledge(
        self, tool_types: list[str] | None = None, cite_ids: list[str] | None = None
    ) -> list[KnowledgeItem]:
        """Get available knowledge (supports filtering)"""
        results = self.knowledge_chain

        if tool_types:
            results = [k for k in results if k.tool_type in tool_types]

        if cite_ids:
            results = [k for k in results if k.cite_id in cite_ids]

        return results

    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,
            "task_id": self.task_id,
            "user_question": self.user_question,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "knowledge_chain": [item.to_dict() for item in self.knowledge_chain],
            "reflections": self.reflections.to_dict(),
            "metadata": self.metadata,
        }
