#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
DR-in-KG 2.0 Core Data Structures
Includes: TopicBlock, ToolTrace, DynamicTopicQueue
"""

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


class TopicStatus(Enum):
    """Topic block status enumeration"""

    PENDING = "pending"  # Pending research
    RESEARCHING = "researching"  # Researching
    COMPLETED = "completed"  # Completed
    FAILED = "failed"  # Failed


class ToolType(Enum):
    """Tool type enumeration"""

    RAG_NAIVE = "rag_naive"
    RAG_HYBRID = "rag_hybrid"
    QUERY_ITEM = "query_item"
    PAPER_SEARCH = "paper_search"
    RUN_CODE = "run_code"
    WEB_SEARCH = "web_search"


# Default max size for raw_answer (50KB)
DEFAULT_RAW_ANSWER_MAX_SIZE = 50 * 1024


@dataclass
class ToolTrace:
    """
    Tool trace - Records complete loop of a single tool call
    """

    tool_id: str  # Unique identifier (e.g., "tool_1", "tool_2")
    citation_id: str  # Citation ID (for report citations and anchors, e.g., CIT-1-01)
    tool_type: str  # Tool type (rag_naive, web_search, etc.)
    query: str  # Query statement issued
    raw_answer: str  # Raw detailed result returned by tool (may be truncated)
    summary: str  # Core summary generated by Note Agent
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    raw_answer_truncated: bool = field(default=False)  # Whether raw_answer was truncated
    raw_answer_original_size: int = field(default=0)  # Original size before truncation

    def __post_init__(self):
        """Post-initialization to handle raw_answer size limit"""
        if self.raw_answer_original_size == 0:
            self.raw_answer_original_size = len(self.raw_answer)

        # Truncate if needed
        if len(self.raw_answer) > DEFAULT_RAW_ANSWER_MAX_SIZE:
            self.raw_answer = self._truncate_raw_answer(
                self.raw_answer, DEFAULT_RAW_ANSWER_MAX_SIZE
            )
            self.raw_answer_truncated = True

    @staticmethod
    def _truncate_raw_answer(raw_answer: str, max_size: int) -> str:
        """
        Truncate raw_answer while trying to preserve valid JSON structure

        Args:
            raw_answer: Original raw answer string
            max_size: Maximum size in bytes

        Returns:
            Truncated string
        """
        if len(raw_answer) <= max_size:
            return raw_answer

        # Try to parse as JSON and truncate intelligently
        try:
            data = json.loads(raw_answer)

            # If it's a dict with common RAG response fields, truncate content fields
            if isinstance(data, dict):
                # Truncate long content fields
                content_fields = ["answer", "content", "text", "chunks", "documents"]
                for field_name in content_fields:
                    if field_name in data:
                        if (
                            isinstance(data[field_name], str)
                            and len(data[field_name]) > max_size // 2
                        ):
                            data[field_name] = data[field_name][: max_size // 2] + "... [truncated]"
                        elif isinstance(data[field_name], list):
                            # Keep only first few items
                            data[field_name] = data[field_name][:3]
                            if data[field_name]:
                                data[field_name].append({"note": "... additional items truncated"})

                truncated = json.dumps(data, ensure_ascii=False)
                if len(truncated) <= max_size:
                    return truncated
        except (json.JSONDecodeError, TypeError):
            pass

        # Fallback: simple truncation with marker
        truncation_marker = "\n... [content truncated, original size: {} bytes]".format(
            len(raw_answer)
        )
        return raw_answer[: max_size - len(truncation_marker)] + truncation_marker

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

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ToolTrace":
        """Create from dictionary"""
        # Handle backward compatibility - old data may not have new fields
        if "raw_answer_truncated" not in data:
            data["raw_answer_truncated"] = False
        if "raw_answer_original_size" not in data:
            data["raw_answer_original_size"] = len(data.get("raw_answer", ""))
        return cls(**data)

    @classmethod
    def create_with_size_limit(
        cls,
        tool_id: str,
        citation_id: str,
        tool_type: str,
        query: str,
        raw_answer: str,
        summary: str,
        max_size: int = DEFAULT_RAW_ANSWER_MAX_SIZE,
    ) -> "ToolTrace":
        """
        Create a ToolTrace with explicit size limit

        Args:
            tool_id: Tool ID
            citation_id: Citation ID
            tool_type: Tool type
            query: Query string
            raw_answer: Raw answer (will be truncated if needed)
            summary: Summary
            max_size: Maximum size for raw_answer

        Returns:
            ToolTrace instance
        """
        original_size = len(raw_answer)
        truncated = len(raw_answer) > max_size

        if truncated:
            raw_answer = cls._truncate_raw_answer(raw_answer, max_size)

        return cls(
            tool_id=tool_id,
            citation_id=citation_id,
            tool_type=tool_type,
            query=query,
            raw_answer=raw_answer,
            summary=summary,
            raw_answer_truncated=truncated,
            raw_answer_original_size=original_size,
        )


@dataclass
class TopicBlock:
    """
    Topic block - Minimum scheduling unit in queue
    """

    block_id: str  # Unique identifier (e.g., "block_1", "block_2")
    sub_topic: str  # Sub-topic name
    overview: str  # Topic overview/background
    status: TopicStatus = TopicStatus.PENDING  # Topic status
    tool_traces: list[ToolTrace] = field(default_factory=list)  # Tool call trace list
    iteration_count: int = 0  # Current iteration count
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
    metadata: dict[str, Any] = field(default_factory=dict)  # Additional metadata

    def add_tool_trace(self, trace: ToolTrace) -> None:
        """Add tool trace"""
        self.tool_traces.append(trace)
        self.updated_at = datetime.now().isoformat()

    def get_latest_trace(self) -> ToolTrace | None:
        """Get latest tool trace"""
        return self.tool_traces[-1] if self.tool_traces else None

    def get_all_summaries(self) -> str:
        """Get concatenated summaries of all tool traces"""
        if not self.tool_traces:
            return ""
        return "\n".join([f"[{trace.tool_type}] {trace.summary}" for trace in self.tool_traces])

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        data = asdict(self)
        data["status"] = self.status.value
        data["tool_traces"] = [trace.to_dict() for trace in self.tool_traces]
        return data

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "TopicBlock":
        """Create from dictionary"""
        data_copy = data.copy()
        if isinstance(data_copy.get("status"), str):
            data_copy["status"] = TopicStatus(data_copy["status"])
        if "tool_traces" in data_copy:
            data_copy["tool_traces"] = [
                ToolTrace.from_dict(t) if isinstance(t, dict) else t
                for t in data_copy["tool_traces"]
            ]
        return cls(**data_copy)


class DynamicTopicQueue:
    """
    Dynamic topic queue - Core memory and scheduling center of the system
    """

    def __init__(
        self, research_id: str, max_length: int | None = None, state_file: str | None = None
    ):
        """
        Initialize queue

        Args:
            research_id: Research task ID
            max_length: Maximum queue length (None means unlimited)
            state_file: Auto-persistence file path
        """
        self.research_id = research_id
        self.blocks: list[TopicBlock] = []
        self.block_counter = 0
        self.created_at = datetime.now().isoformat()
        self.max_length = max_length if isinstance(max_length, int) and max_length > 0 else None
        self.state_file = state_file

    def set_state_file(self, filepath: str | None) -> None:
        """Set queue auto-persistence file"""
        self.state_file = filepath
        self._auto_save()

    @staticmethod
    def _normalize_topic(text: str) -> str:
        return (text or "").strip().lower()

    def add_block(self, sub_topic: str, overview: str) -> TopicBlock:
        """
        Add new topic block to the end of queue

        Args:
            sub_topic: Sub-topic name
            overview: Topic overview

        Returns:
            Created TopicBlock
        """
        if self.max_length and len(self.blocks) >= self.max_length:
            raise RuntimeError(
                f"Queue has reached maximum capacity ({self.max_length}), cannot add new topic."
            )
        self.block_counter += 1
        block_id = f"block_{self.block_counter}"
        block = TopicBlock(block_id=block_id, sub_topic=sub_topic, overview=overview)
        self.blocks.append(block)
        self._auto_save()
        return block

    def has_topic(self, sub_topic: str) -> bool:
        """Check if topic already exists (case-insensitive, ignoring leading/trailing spaces)"""
        target = self._normalize_topic(sub_topic)
        if not target:
            return False
        return any(self._normalize_topic(b.sub_topic) == target for b in self.blocks)

    def list_topics(self) -> list[str]:
        """List all current topic titles"""
        return [b.sub_topic for b in self.blocks]

    def get_pending_block(self) -> TopicBlock | None:
        """
        Get first pending topic block

        Returns:
            First TopicBlock with PENDING status, or None if not found
        """
        for block in self.blocks:
            if block.status == TopicStatus.PENDING:
                return block
        return None

    def get_block_by_id(self, block_id: str) -> TopicBlock | None:
        """
        Get topic block by ID

        Args:
            block_id: Topic block ID

        Returns:
            Corresponding TopicBlock, or None if not found
        """
        for block in self.blocks:
            if block.block_id == block_id:
                return block
        return None

    def mark_researching(self, block_id: str) -> bool:
        """
        Mark topic block as researching

        Args:
            block_id: Topic block ID

        Returns:
            Whether marking was successful
        """
        block = self.get_block_by_id(block_id)
        if block:
            block.status = TopicStatus.RESEARCHING
            block.updated_at = datetime.now().isoformat()
            self._auto_save()
            return True
        return False

    def mark_completed(self, block_id: str) -> bool:
        """
        Mark topic block as completed

        Args:
            block_id: Topic block ID

        Returns:
            Whether marking was successful
        """
        block = self.get_block_by_id(block_id)
        if block:
            block.status = TopicStatus.COMPLETED
            block.updated_at = datetime.now().isoformat()
            self._auto_save()
            return True
        return False

    def mark_failed(self, block_id: str) -> bool:
        """
        Mark topic block as failed

        Args:
            block_id: Topic block ID

        Returns:
            Whether marking was successful
        """
        block = self.get_block_by_id(block_id)
        if block:
            block.status = TopicStatus.FAILED
            block.updated_at = datetime.now().isoformat()
            self._auto_save()
            return True
        return False

    def get_all_completed_blocks(self) -> list[TopicBlock]:
        """Get all completed topic blocks"""
        return [b for b in self.blocks if b.status == TopicStatus.COMPLETED]

    def get_all_pending_blocks(self) -> list[TopicBlock]:
        """Get all pending topic blocks"""
        return [b for b in self.blocks if b.status == TopicStatus.PENDING]

    def is_all_completed(self) -> bool:
        """Check if all topic blocks are completed"""
        if not self.blocks:
            return False
        return all(b.status == TopicStatus.COMPLETED for b in self.blocks)

    def get_statistics(self) -> dict[str, Any]:
        """Get queue statistics"""
        return {
            "total_blocks": len(self.blocks),
            "pending": len(self.get_all_pending_blocks()),
            "researching": len([b for b in self.blocks if b.status == TopicStatus.RESEARCHING]),
            "completed": len(self.get_all_completed_blocks()),
            "failed": len([b for b in self.blocks if b.status == TopicStatus.FAILED]),
            "total_tool_calls": sum(len(b.tool_traces) for b in self.blocks),
        }

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        return {
            "research_id": self.research_id,
            "created_at": self.created_at,
            "blocks": [b.to_dict() for b in self.blocks],
            "statistics": self.get_statistics(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DynamicTopicQueue":
        """Create from dictionary"""
        queue = cls(data["research_id"])
        queue.created_at = data.get("created_at", queue.created_at)
        for block_data in data.get("blocks", []):
            block = TopicBlock.from_dict(block_data)
            queue.blocks.append(block)
            # Update counter
            if block.block_id.startswith("block_"):
                try:
                    block_num = int(block.block_id.split("_")[1])
                    queue.block_counter = max(queue.block_counter, block_num)
                except (ValueError, IndexError):
                    pass
        return queue

    def save_to_json(self, filepath: str) -> None:
        """Save queue to JSON file"""
        Path(filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, ensure_ascii=False, indent=2)

    def _auto_save(self) -> None:
        """Auto-save if state_file is set"""
        if self.state_file:
            try:
                self.save_to_json(self.state_file)
            except Exception as exc:
                print(f"⚠️ Failed to save queue progress: {exc}")

    @classmethod
    def load_from_json(cls, filepath: str) -> "DynamicTopicQueue":
        """Load queue from JSON file"""
        with open(filepath, encoding="utf-8") as f:
            data = json.load(f)
        return cls.from_dict(data)


__all__ = [
    "DynamicTopicQueue",
    "ToolTrace",
    "ToolType",
    "TopicBlock",
    "TopicStatus",
]
