"""
MCAL Core Data Models

Defines the fundamental data structures for:
- Intent Graphs (goal hierarchies)
- Decision Trails (reasoning preservation)
- Memory objects (context storage)
"""

from __future__ import annotations

from datetime import datetime, timezone
from enum import Enum
from typing import Optional
from uuid import uuid4

from pydantic import BaseModel, Field


def _utc_now() -> datetime:
    """Return current UTC time (timezone-aware)."""
    return datetime.now(timezone.utc)


# =============================================================================
# Enums
# =============================================================================

class IntentType(str, Enum):
    """Types of intent nodes in the hierarchy."""
    MISSION = "mission"      # Overarching objective (rarely changes)
    GOAL = "goal"            # Major sub-objectives (session-level)
    TASK = "task"            # Concrete actions (turn-level)
    DECISION = "decision"    # Choices made with rationale


class IntentStatus(str, Enum):
    """Status of an intent node."""
    ACTIVE = "active"
    COMPLETED = "completed"
    ABANDONED = "abandoned"
    PENDING = "pending"
    BLOCKED = "blocked"


class EdgeRelation(str, Enum):
    """Types of relationships between intent nodes."""
    DERIVES_FROM = "derives_from"    # Child derives from parent
    ENABLES = "enables"              # Completing this enables another
    CONFLICTS_WITH = "conflicts_with"  # Mutually exclusive
    SUPERSEDES = "supersedes"        # Replaces a previous intent
    DEPENDS_ON = "depends_on"        # Requires another to be completed first


class EvidenceSource(str, Enum):
    """Source of evidence supporting a decision."""
    USER_STATED = "user_stated"      # Explicitly stated by user
    INFERRED = "inferred"            # Inferred from context
    EXTERNAL = "external"            # From external source (web, docs)
    SYSTEM = "system"                # System-generated


# =============================================================================
# Intent Graph Models
# =============================================================================

class IntentNode(BaseModel):
    """
    A node in the intent graph representing a goal, task, or decision.
    
    Example:
        IntentNode(
            type=IntentType.GOAL,
            content="Build fraud detection ML pipeline",
            status=IntentStatus.ACTIVE,
            confidence=0.9,
            evidence=["turn_5", "turn_8"]
        )
    """
    id: str = Field(default_factory=lambda: str(uuid4())[:8])
    type: IntentType
    content: str
    status: IntentStatus = IntentStatus.ACTIVE
    confidence: float = Field(ge=0.0, le=1.0, default=0.8)
    evidence: list[str] = Field(default_factory=list)  # Turn IDs supporting this
    created_at: datetime = Field(default_factory=_utc_now)
    updated_at: datetime = Field(default_factory=_utc_now)
    metadata: dict = Field(default_factory=dict)
    
    def update_status(self, new_status: IntentStatus) -> None:
        """Update status and timestamp."""
        self.status = new_status
        self.updated_at = _utc_now()


class IntentEdge(BaseModel):
    """
    An edge connecting two intent nodes.
    
    Example:
        IntentEdge(
            source="goal_123",
            target="task_456",
            relation=EdgeRelation.DERIVES_FROM
        )
    """
    id: str = Field(default_factory=lambda: str(uuid4())[:8])
    source: str  # Source node ID
    target: str  # Target node ID
    relation: EdgeRelation
    strength: float = Field(ge=0.0, le=1.0, default=1.0)
    created_at: datetime = Field(default_factory=_utc_now)


class IntentGraph(BaseModel):
    """
    Hierarchical representation of user goals and their relationships.
    
    The graph captures:
    - What the user is trying to achieve (goals)
    - How goals break down into tasks
    - Dependencies and conflicts between goals
    - Progress status of each node
    """
    id: str = Field(default_factory=lambda: str(uuid4())[:8])
    session_id: Optional[str] = None
    nodes: dict[str, IntentNode] = Field(default_factory=dict)
    edges: list[IntentEdge] = Field(default_factory=list)
    root_node_id: Optional[str] = None  # Mission-level node
    created_at: datetime = Field(default_factory=_utc_now)
    updated_at: datetime = Field(default_factory=_utc_now)
    
    def add_node(self, node: IntentNode) -> str:
        """Add a node to the graph."""
        self.nodes[node.id] = node
        self.updated_at = _utc_now()
        if node.type == IntentType.MISSION and self.root_node_id is None:
            self.root_node_id = node.id
        return node.id
    
    def add_edge(self, edge: IntentEdge) -> str:
        """Add an edge to the graph."""
        self.edges.append(edge)
        self.updated_at = _utc_now()
        return edge.id
    
    def get_active_goals(self) -> list[IntentNode]:
        """Return all currently active goals, missions, and tasks.
        
        Includes nodes that are:
        - ACTIVE status (explicitly in progress)
        - PENDING status (not yet started but still active work)
        
        And are of type:
        - MISSION (high-level objectives)
        - GOAL (concrete goals)
        - TASK (actionable tasks)
        """
        active_statuses = (IntentStatus.ACTIVE, IntentStatus.PENDING)
        active_types = (IntentType.MISSION, IntentType.GOAL, IntentType.TASK)
        
        return [
            node for node in self.nodes.values()
            if node.status in active_statuses
            and node.type in active_types
        ]
    
    def get_children(self, node_id: str) -> list[IntentNode]:
        """Get all child nodes of a given node."""
        child_ids = [
            edge.target for edge in self.edges
            if edge.source == node_id and edge.relation == EdgeRelation.DERIVES_FROM
        ]
        return [self.nodes[cid] for cid in child_ids if cid in self.nodes]
    
    def get_node_path(self, node_id: str) -> list[IntentNode]:
        """Get path from root to a specific node."""
        path = []
        current_id = node_id
        visited = set()
        
        while current_id and current_id not in visited:
            visited.add(current_id)
            if current_id in self.nodes:
                path.append(self.nodes[current_id])
            
            # Find parent
            parent_edge = next(
                (e for e in self.edges 
                 if e.target == current_id and e.relation == EdgeRelation.DERIVES_FROM),
                None
            )
            current_id = parent_edge.source if parent_edge else None
        
        return list(reversed(path))


# =============================================================================
# Decision Trail Models
# =============================================================================

class Alternative(BaseModel):
    """An alternative that was considered but not chosen."""
    option: str
    pros: list[str] = Field(default_factory=list)
    cons: list[str] = Field(default_factory=list)
    rejection_reason: Optional[str] = None


class Evidence(BaseModel):
    """Evidence supporting a decision."""
    claim: str
    source: EvidenceSource
    confidence: float = Field(ge=0.0, le=1.0, default=0.8)
    turn_id: Optional[str] = None
    external_url: Optional[str] = None


class TradeOff(BaseModel):
    """A trade-off acknowledged in making a decision."""
    gained: str
    sacrificed: str
    justification: Optional[str] = None


class DecisionTrail(BaseModel):
    """
    Captures not just WHAT was decided, but WHY.
    
    This is the core innovation for reasoning preservation:
    instead of storing "User chose PostgreSQL", we store:
    - The decision itself
    - What alternatives were considered
    - Why this option was chosen
    - What evidence supported it
    - What trade-offs were made
    
    Example:
        DecisionTrail(
            decision="Use PostgreSQL for the database",
            context="Discussing data storage for ML pipeline",
            alternatives=[
                Alternative(option="MongoDB", rejection_reason="Need ACID compliance"),
                Alternative(option="SQLite", rejection_reason="Won't scale")
            ],
            rationale="PostgreSQL offers ACID + pgvector for embeddings",
            evidence=[Evidence(claim="pgvector supports vector similarity", source=EvidenceSource.EXTERNAL)]
        )
    """
    id: str = Field(default_factory=lambda: str(uuid4())[:8])
    decision: str  # What was decided
    context: str   # Situation when decision was made
    
    # The "WHY" components
    alternatives: list[Alternative] = Field(default_factory=list)
    rationale: str  # Why this option was chosen
    evidence: list[Evidence] = Field(default_factory=list)
    trade_offs: list[TradeOff] = Field(default_factory=list)
    confidence: float = Field(ge=0.0, le=1.0, default=0.8)
    
    # Linkage
    related_goals: list[str] = Field(default_factory=list)  # Goal IDs this serves
    dependencies: list[str] = Field(default_factory=list)   # Prior decisions this builds on
    invalidated_by: Optional[str] = None  # If superseded, what replaced it
    
    # Metadata
    turn_id: Optional[str] = None
    created_at: datetime = Field(default_factory=_utc_now)
    updated_at: datetime = Field(default_factory=_utc_now)
    
    def invalidate(self, replacement_id: str) -> None:
        """Mark this decision as invalidated by a new decision."""
        self.invalidated_by = replacement_id
        self.updated_at = _utc_now()
    
    @property
    def is_valid(self) -> bool:
        """Check if this decision is still valid (not superseded)."""
        return self.invalidated_by is None


# =============================================================================
# Memory Models
# =============================================================================

class MemoryType(str, Enum):
    """Types of memory items."""
    FACT = "fact"            # Simple factual information
    PREFERENCE = "preference"  # User preference
    INTENT = "intent"        # Goal/intent (links to IntentNode)
    DECISION = "decision"    # Decision (links to DecisionTrail)
    EPISODE = "episode"      # Episodic memory (conversation summary)


class Memory(BaseModel):
    """
    A memory item that can be stored and retrieved.
    
    This is the base unit of storage, which can represent:
    - Simple facts (like Mem0)
    - Intent graph nodes
    - Decision trails
    - Episodic summaries
    """
    id: str = Field(default_factory=lambda: str(uuid4())[:8])
    type: MemoryType
    content: str
    
    # Embedding for semantic search
    embedding: Optional[list[float]] = None
    
    # Relevance factors
    importance: float = Field(ge=0.0, le=1.0, default=0.5)
    user_marked: bool = False  # Explicitly marked important by user
    reference_count: int = 0   # How often this has been referenced
    
    # Linkage to structured data
    intent_node_id: Optional[str] = None
    decision_trail_id: Optional[str] = None
    
    # Metadata
    session_id: Optional[str] = None
    turn_id: Optional[str] = None
    created_at: datetime = Field(default_factory=_utc_now)
    last_accessed: datetime = Field(default_factory=_utc_now)
    
    def access(self) -> None:
        """Record that this memory was accessed."""
        self.reference_count += 1
        self.last_accessed = _utc_now()


# =============================================================================
# Conversation Models
# =============================================================================

class Session(BaseModel):
    """A conversation session."""
    id: str = Field(default_factory=lambda: str(uuid4())[:8])
    intent_graph: Optional[IntentGraph] = None
    decisions: list[DecisionTrail] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=_utc_now)
    updated_at: datetime = Field(default_factory=_utc_now)


# =============================================================================
# Retrieval Models
# =============================================================================

class RetrievalResult(BaseModel):
    """Result of a memory retrieval operation."""
    memory: Memory
    score: float  # Combined relevance score
    score_breakdown: dict[str, float] = Field(default_factory=dict)  # Component scores
    
    
class RetrievalConfig(BaseModel):
    """Configuration for retrieval operations."""
    max_results: int = 10
    min_score: float = 0.0
    
    # Weight factors for multi-factor scoring
    semantic_weight: float = 0.4
    goal_alignment_weight: float = 0.3
    recency_weight: float = 0.1
    reference_weight: float = 0.1
    decision_impact_weight: float = 0.1
    
    # Filters
    memory_types: Optional[list[MemoryType]] = None
    include_invalidated_decisions: bool = False
