"""
H-MEM Level 1: Trace dataclass and TraceConsolidator.

Traces consolidate multiple episodes into coherent narratives with
themes, summaries, and extracted patterns for higher-level memory.
"""

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import List, Optional, Dict, Any, TYPE_CHECKING
import uuid

from gsd_rlm.memory.hmem.episode import Episode, EpisodeType

if TYPE_CHECKING:
    from gsd_rlm.memory.hmem.store import EpisodeStore


def _utc_now_iso() -> str:
    """Get current UTC time as ISO format string."""
    return datetime.now(timezone.utc).isoformat()


def _generate_trace_id() -> str:
    """Generate a unique trace ID."""
    return f"trace_{uuid.uuid4().hex[:12]}"


@dataclass
class Trace:
    """
    A consolidated narrative from multiple episodes (H-MEM Level 1).

    Traces represent coherent sequences of agent actions that form
    a meaningful narrative, with extracted themes and lessons learned.

    This is the first level of memory abstraction above episodes.
    """

    # Required identification fields
    trace_id: str
    agent_id: str
    session_id: str
    theme: str  # Main theme/topic of this trace

    # Episode linkage
    episode_ids: List[str] = field(default_factory=list)
    episode_count: int = 0

    # Consolidated content
    summary: str = ""
    patterns_identified: List[str] = field(default_factory=list)
    overall_success: bool = False
    lessons_learned: List[str] = field(default_factory=list)

    # Hierarchy linkage
    category_id: Optional[str] = None  # Links to Level 2 Category

    # Metadata
    created_at: str = field(default_factory=_utc_now_iso)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert trace to dictionary for serialization.

        Returns:
            Dictionary representation suitable for JSON storage.
        """
        return {
            "trace_id": self.trace_id,
            "agent_id": self.agent_id,
            "session_id": self.session_id,
            "theme": self.theme,
            "episode_ids": self.episode_ids,
            "episode_count": self.episode_count,
            "summary": self.summary,
            "patterns_identified": self.patterns_identified,
            "overall_success": self.overall_success,
            "lessons_learned": self.lessons_learned,
            "category_id": self.category_id,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Trace":
        """
        Create Trace from dictionary.

        Args:
            data: Dictionary with trace fields.

        Returns:
            Trace instance.
        """
        return cls(
            trace_id=data.get("trace_id", _generate_trace_id()),
            agent_id=data.get("agent_id", ""),
            session_id=data.get("session_id", ""),
            theme=data.get("theme", ""),
            episode_ids=data.get("episode_ids", []),
            episode_count=data.get("episode_count", 0),
            summary=data.get("summary", ""),
            patterns_identified=data.get("patterns_identified", []),
            overall_success=data.get("overall_success", False),
            lessons_learned=data.get("lessons_learned", []),
            category_id=data.get("category_id"),
            created_at=data.get("created_at", _utc_now_iso()),
        )

    def __post_init__(self):
        """Validate and normalize fields after initialization."""
        # Ensure lists are always lists
        if self.episode_ids is None:
            self.episode_ids = []
        if self.patterns_identified is None:
            self.patterns_identified = []
        if self.lessons_learned is None:
            self.lessons_learned = []

        # Ensure episode_count matches list length
        if self.episode_count == 0 and self.episode_ids:
            self.episode_count = len(self.episode_ids)

        # Ensure created_at has a value
        if not self.created_at:
            self.created_at = _utc_now_iso()

    @property
    def is_categorized(self) -> bool:
        """Check if trace has been categorized."""
        return self.category_id is not None


class TraceConsolidator:
    """
    Consolidates episodes into traces (H-MEM Level 1).

    Uses LLM for intelligent summarization when available,
    falls back to simple aggregation otherwise.
    """

    def __init__(
        self,
        episode_store: "EpisodeStore",
        llm_provider: Optional[Any] = None,
    ):
        """
        Initialize the trace consolidator.

        Args:
            episode_store: Store to retrieve episodes from.
            llm_provider: Optional LLM provider for summarization.
                         If None, uses simple aggregation.
        """
        self.episode_store = episode_store
        self.llm_provider = llm_provider

    async def consolidate_session(self, session_id: str) -> Optional[Trace]:
        """
        Consolidate all episodes from a session into a single trace.

        Args:
            session_id: Session identifier to consolidate.

        Returns:
            Created Trace if episodes exist, None otherwise.
        """
        # Get all episodes for this session
        episodes = self.episode_store.get_episodes_for_session(session_id)

        if not episodes:
            return None

        # Filter to only unconsolidated episodes
        unconsolidated = [e for e in episodes if not e.is_consolidated]

        if not unconsolidated:
            return None

        # Create trace from episodes
        trace = await self._create_trace_from_episodes(unconsolidated)

        # Update episode trace_ids
        for episode in unconsolidated:
            episode.mark_consolidated(trace.trace_id)
            self.episode_store.store(episode)

        return trace

    async def _create_trace_from_episodes(self, episodes: List[Episode]) -> Trace:
        """
        Create a trace from a list of episodes.

        Args:
            episodes: Episodes to consolidate.

        Returns:
            Consolidated Trace.
        """
        if not episodes:
            raise ValueError("Cannot create trace from empty episode list")

        # Use first episode for agent/session info
        first_episode = episodes[0]

        # Build trace with basic info
        trace = Trace(
            trace_id=_generate_trace_id(),
            agent_id=first_episode.agent_id,
            session_id=first_episode.session_id,
            episode_ids=[e.episode_id for e in episodes],
            episode_count=len(episodes),
            theme="",  # Will be set by consolidation
        )

        # Try LLM consolidation if available
        if self.llm_provider is not None:
            try:
                consolidation = await self._consolidate_with_llm(episodes)
                trace.theme = consolidation.get("theme", "")
                trace.summary = consolidation.get("summary", "")
                trace.patterns_identified = consolidation.get("patterns", [])
                trace.lessons_learned = consolidation.get("lessons", [])
                trace.overall_success = consolidation.get("overall_success", False)
            except Exception:
                # Fall back to simple aggregation on LLM error
                self._simple_aggregation(trace, episodes)
        else:
            # No LLM, use simple aggregation
            self._simple_aggregation(trace, episodes)

        return trace

    def _build_consolidation_prompt(self, episodes: List[Episode]) -> str:
        """
        Build a prompt for LLM consolidation.

        Args:
            episodes: Episodes to format.

        Returns:
            Formatted prompt string.
        """
        episode_summaries = []
        for i, ep in enumerate(episodes, 1):
            episode_summaries.append(
                f"Episode {i} ({ep.episode_type.value}):\n"
                f"  Context: {ep.context[:200]}...\n"
                f"  Action: {ep.action[:200]}...\n"
                f"  Outcome: {ep.outcome[:200]}...\n"
                f"  Success: {ep.success}"
            )

        episodes_text = "\n\n".join(episode_summaries)

        return f"""Analyze the following agent episodes and consolidate them into a coherent trace.

Episodes:
{episodes_text}

Provide a JSON response with:
1. "theme": A short (5-10 words) main theme of this sequence
2. "summary": A 2-3 sentence summary of what happened
3. "patterns": List of 2-5 patterns identified in the episodes
4. "lessons": List of 1-3 lessons learned
5. "overall_success": Boolean indicating if the overall sequence was successful

Response as valid JSON only."""

    def _parse_consolidation_response(self, response: str) -> dict:
        """
        Parse LLM response into consolidation data.

        Args:
            response: Raw LLM response string.

        Returns:
            Parsed consolidation dictionary.
        """
        import json

        # Try to extract JSON from response
        try:
            # Direct JSON parse
            return json.loads(response)
        except json.JSONDecodeError:
            pass

        # Try to find JSON in response
        import re

        json_match = re.search(r"\{[\s\S]*\}", response)
        if json_match:
            try:
                return json.loads(json_match.group())
            except json.JSONDecodeError:
                pass

        # Return defaults if parsing fails
        return {
            "theme": "Agent activity",
            "summary": "Agent completed multiple episodes.",
            "patterns": [],
            "lessons": [],
            "overall_success": False,
        }

    async def _consolidate_with_llm(self, episodes: List[Episode]) -> dict:
        """
        Use LLM to consolidate episodes.

        Args:
            episodes: Episodes to consolidate.

        Returns:
            Consolidation data dictionary.
        """
        prompt = self._build_consolidation_prompt(episodes)

        # Call LLM provider - handle different provider interfaces
        if hasattr(self.llm_provider, "generate"):
            response = await self.llm_provider.generate(prompt)
        elif hasattr(self.llm_provider, "complete"):
            response = await self.llm_provider.complete(prompt)
        elif hasattr(self.llm_provider, "chat"):
            response = await self.llm_provider.chat(prompt)
        elif callable(self.llm_provider):
            response = self.llm_provider(prompt)
        else:
            raise ValueError("LLM provider has no recognized method")

        return self._parse_consolidation_response(response)

    def _simple_aggregation(self, trace: Trace, episodes: List[Episode]) -> None:
        """
        Simple aggregation without LLM.

        Args:
            trace: Trace to populate.
            episodes: Episodes to aggregate.
        """
        # Determine theme from episode types
        episode_types = [e.episode_type.value for e in episodes]
        most_common_type = max(set(episode_types), key=episode_types.count)
        trace.theme = f"{most_common_type.replace('_', ' ').title()} session"

        # Build summary from episode outcomes
        outcomes = [e.outcome for e in episodes if e.outcome]
        if outcomes:
            trace.summary = " ".join(outcomes[:3])[:500]
        else:
            trace.summary = f"Session with {len(episodes)} episodes."

        # Extract simple patterns from tags
        all_tags = []
        for ep in episodes:
            all_tags.extend(ep.tags)
        if all_tags:
            # Get unique tags as patterns
            trace.patterns_identified = list(set(all_tags))[:5]

        # Determine overall success
        successes = [e.success for e in episodes]
        trace.overall_success = all(successes) if successes else False

        # Simple lessons based on failures
        failed_episodes = [e for e in episodes if not e.success]
        if failed_episodes:
            trace.lessons_learned = [
                f"Encountered {len(failed_episodes)} issue(s) requiring attention"
            ]
