"""
Intent Tracker

Extracts and maintains hierarchical intent graphs from conversations.
This is Pillar 1 of MCAL: Intent Graph Preservation.

Key capabilities:
- Extract intent structures from conversation turns
- Incrementally update intent graph as conversation progresses
- Track goal status (active, completed, abandoned)
- Detect intent drift and evolution
"""

from __future__ import annotations

import json
import logging
from typing import Optional, Protocol

from .models import (
    EdgeRelation,
    IntentEdge,
    IntentGraph,
    IntentNode,
    IntentStatus,
    IntentType,
    Turn,
)

logger = logging.getLogger(__name__)


# Issue #1: Map invalid intent types from LLM to valid enum values
INTENT_TYPE_MAPPING = {
    "evidence": IntentType.TASK,
    "consideration": IntentType.TASK,
    "analysis": IntentType.TASK,
    "objective": IntentType.GOAL,
    "sub-goal": IntentType.TASK,
    "question": IntentType.TASK,
    "constraint": IntentType.TASK,
    "preference": IntentType.TASK,
}

# Issue #4: Map invalid intent status from LLM to valid enum values
INTENT_STATUS_MAPPING = {
    "reopened": IntentStatus.ACTIVE,
    "paused": IntentStatus.BLOCKED,
    "in_progress": IntentStatus.ACTIVE,
    "in-progress": IntentStatus.ACTIVE,
    "done": IntentStatus.COMPLETED,
    "cancelled": IntentStatus.ABANDONED,
    "canceled": IntentStatus.ABANDONED,
}

# Issue #5: Map invalid edge relations from LLM to valid enum values
EDGE_RELATION_MAPPING = {
    "results_in": EdgeRelation.ENABLES,
    "leads_to": EdgeRelation.ENABLES,
    "requires": EdgeRelation.DEPENDS_ON,
    "needs": EdgeRelation.DEPENDS_ON,
    "parent_of": EdgeRelation.DERIVES_FROM,
    "child_of": EdgeRelation.DERIVES_FROM,
    "replaces": EdgeRelation.SUPERSEDES,
    "blocks": EdgeRelation.CONFLICTS_WITH,
}


def normalize_intent_type(type_str: str) -> IntentType:
    """
    Normalize intent type string to valid IntentType enum.
    
    Handles invalid types from LLM by mapping to closest valid type.
    Fixes Issue #1: Invalid IntentType values from LLM extraction.
    
    Args:
        type_str: Raw type string from LLM
        
    Returns:
        Valid IntentType enum value
    """
    type_lower = type_str.lower().strip()
    
    # Try direct enum conversion first
    try:
        return IntentType(type_lower)
    except ValueError:
        pass
    
    # Try mapping table
    if type_lower in INTENT_TYPE_MAPPING:
        logger.warning(f"Mapped invalid intent type '{type_str}' to {INTENT_TYPE_MAPPING[type_lower]}")
        return INTENT_TYPE_MAPPING[type_lower]
    
    # Fallback to TASK for unknown types
    logger.warning(f"Unknown intent type '{type_str}', falling back to TASK")
    return IntentType.TASK


def normalize_intent_status(status_str: str) -> IntentStatus:
    """
    Normalize intent status string to valid IntentStatus enum.
    
    Fixes Issue #4: Invalid IntentStatus values from LLM.
    """
    status_lower = status_str.lower().strip()
    
    try:
        return IntentStatus(status_lower)
    except ValueError:
        pass
    
    if status_lower in INTENT_STATUS_MAPPING:
        logger.warning(f"Mapped invalid intent status '{status_str}' to {INTENT_STATUS_MAPPING[status_lower]}")
        return INTENT_STATUS_MAPPING[status_lower]
    
    logger.warning(f"Unknown intent status '{status_str}', falling back to ACTIVE")
    return IntentStatus.ACTIVE


def normalize_edge_relation(relation_str: str) -> EdgeRelation:
    """
    Normalize edge relation string to valid EdgeRelation enum.
    
    Fixes Issue #5: Invalid EdgeRelation values from LLM.
    """
    relation_lower = relation_str.lower().strip()
    
    try:
        return EdgeRelation(relation_lower)
    except ValueError:
        pass
    
    if relation_lower in EDGE_RELATION_MAPPING:
        logger.warning(f"Mapped invalid edge relation '{relation_str}' to {EDGE_RELATION_MAPPING[relation_lower]}")
        return EDGE_RELATION_MAPPING[relation_lower]
    
    logger.warning(f"Unknown edge relation '{relation_str}', falling back to DERIVES_FROM")
    return EdgeRelation.DERIVES_FROM


# =============================================================================
# LLM Client Protocol
# =============================================================================

class LLMClient(Protocol):
    """Protocol for LLM client implementations."""
    
    async def complete(self, prompt: str, system: Optional[str] = None) -> str:
        """Generate a completion for the given prompt."""
        ...


# =============================================================================
# Prompts
# =============================================================================

INTENT_EXTRACTION_SYSTEM = """You are an expert at analyzing conversations to extract user intents and goals.

Your task is to identify the hierarchical structure of what the user is trying to achieve:
- MISSION: The overarching objective (if discernible)
- GOALS: Major sub-objectives the user wants to accomplish
- TASKS: Specific actions or steps within goals
- DECISIONS: Choices the user has made or needs to make

For each intent, assess:
- Status: active, completed, abandoned, pending, or blocked
- Confidence: How certain are you this is a real intent (0.0-1.0)
- Evidence: Which parts of the conversation support this intent

Output your analysis as valid JSON."""

INTENT_EXTRACTION_PROMPT = """Analyze this conversation and extract the user's intent hierarchy.

CONVERSATION:
{conversation}

Extract intents as JSON with this structure:
{{
    "mission": {{
        "content": "string or null if not clear",
        "confidence": 0.0-1.0
    }},
    "goals": [
        {{
            "id": "g1",
            "content": "goal description",
            "status": "active|completed|abandoned|pending|blocked",
            "confidence": 0.0-1.0,
            "evidence": ["turn_1", "turn_3"],
            "parent": "mission or null"
        }}
    ],
    "tasks": [
        {{
            "id": "t1", 
            "content": "task description",
            "status": "active|completed|abandoned|pending|blocked",
            "confidence": 0.0-1.0,
            "evidence": ["turn_2"],
            "parent": "g1"
        }}
    ],
    "decisions": [
        {{
            "id": "d1",
            "content": "decision description",
            "status": "active|completed",
            "confidence": 0.0-1.0,
            "evidence": ["turn_4"],
            "parent": "t1 or g1"
        }}
    ]
}}

Be thorough but don't invent intents that aren't supported by the conversation.
Output ONLY valid JSON, no explanation."""

INTENT_UPDATE_PROMPT = """Given the existing intent graph and a new conversation turn, update the intent structure.

EXISTING INTENT GRAPH:
{intent_graph}

NEW TURN:
{turn}

Determine what changes are needed:
1. New intents to add?
2. Status changes for existing intents?
3. New relationships between intents?

Output as JSON:
{{
    "new_nodes": [
        {{
            "id": "unique_id",
            "type": "mission|goal|task|decision",
            "content": "Description of the intent",
            "status": "active|completed|pending|abandoned|blocked",
            "confidence": 0.0-1.0,
            "parent": "parent_node_id or null"
        }}
    ],
    "status_updates": [
        {{"id": "g1", "new_status": "completed"}}
    ],
    "new_edges": [
        {{"source": "g1", "target": "t2", "relation": "derives_from"}}
    ]
}}

IMPORTANT: Each new_node MUST have a "content" field with a description.
Output ONLY valid JSON, no explanation."""


# =============================================================================
# Intent Tracker
# =============================================================================

class IntentTracker:
    """
    Extracts and maintains intent graphs from conversations.
    
    Usage:
        tracker = IntentTracker(llm_client)
        
        # Extract from full conversation
        graph = await tracker.extract_intents(turns)
        
        # Or incrementally update
        graph = await tracker.update_intent(new_turn, existing_graph)
        
        # Query active goals
        active = tracker.get_active_goals(graph)
    """
    
    def __init__(self, llm_client: LLMClient):
        """
        Initialize the intent tracker.
        
        Args:
            llm_client: LLM client for extraction (Anthropic, OpenAI, etc.)
        """
        self.llm = llm_client
    
    async def extract_intents(
        self,
        turns: list[Turn],
        session_id: Optional[str] = None
    ) -> IntentGraph:
        """
        Extract intent graph from a conversation.
        
        Args:
            turns: List of conversation turns
            session_id: Optional session identifier
            
        Returns:
            IntentGraph with extracted intents
        """
        # Format conversation for prompt (use smart formatting for long conversations)
        conversation = await self._format_conversation_smart(turns)
        
        # Call LLM for extraction
        prompt = INTENT_EXTRACTION_PROMPT.format(conversation=conversation)
        response = await self.llm.complete(prompt, system=INTENT_EXTRACTION_SYSTEM)
        
        # Parse response
        try:
            data = json.loads(self._clean_json_response(response))
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse intent extraction response: {e}")
            logger.debug(f"Response was: {response}")
            return IntentGraph(session_id=session_id)
        
        # Build graph
        graph = IntentGraph(session_id=session_id)
        node_id_map: dict[str, str] = {}  # Map from LLM IDs to our IDs
        
        # Add mission if present
        mission_data = data.get("mission", {})
        mission_content = self._extract_node_content(mission_data) if mission_data else None
        if mission_content:
            mission_node = IntentNode(
                type=IntentType.MISSION,
                content=mission_content,
                status=IntentStatus.ACTIVE,
                confidence=mission_data.get("confidence", 0.8)
            )
            graph.add_node(mission_node)
            node_id_map["mission"] = mission_node.id
        
        # Add goals
        for goal_data in data.get("goals", []):
            content = self._extract_node_content(goal_data)
            if not content:
                continue
            node = IntentNode(
                type=IntentType.GOAL,
                content=content,
                status=normalize_intent_status(goal_data.get("status", "active")),
                confidence=goal_data.get("confidence", 0.8),
                evidence=goal_data.get("evidence", [])
            )
            graph.add_node(node)
            node_id_map[goal_data["id"]] = node.id
            
            # Add edge to parent
            parent_key = goal_data.get("parent")
            if parent_key and parent_key in node_id_map:
                edge = IntentEdge(
                    source=node_id_map[parent_key],
                    target=node.id,
                    relation=EdgeRelation.DERIVES_FROM
                )
                graph.add_edge(edge)
        
        # Add tasks
        for task_data in data.get("tasks", []):
            content = self._extract_node_content(task_data)
            if not content:
                continue
            node = IntentNode(
                type=IntentType.TASK,
                content=content,
                status=normalize_intent_status(task_data.get("status", "active")),
                confidence=task_data.get("confidence", 0.8),
                evidence=task_data.get("evidence", [])
            )
            graph.add_node(node)
            node_id_map[task_data["id"]] = node.id
            
            # Add edge to parent
            parent_key = task_data.get("parent")
            if parent_key and parent_key in node_id_map:
                edge = IntentEdge(
                    source=node_id_map[parent_key],
                    target=node.id,
                    relation=EdgeRelation.DERIVES_FROM
                )
                graph.add_edge(edge)
        
        # Add decisions
        for decision_data in data.get("decisions", []):
            content = self._extract_node_content(decision_data)
            if not content:
                continue
            node = IntentNode(
                type=IntentType.DECISION,
                content=content,
                status=normalize_intent_status(decision_data.get("status", "active")),
                confidence=decision_data.get("confidence", 0.8),
                evidence=decision_data.get("evidence", [])
            )
            graph.add_node(node)
            node_id_map[decision_data["id"]] = node.id
            
            # Add edge to parent
            parent_key = decision_data.get("parent")
            if parent_key and parent_key in node_id_map:
                edge = IntentEdge(
                    source=node_id_map[parent_key],
                    target=node.id,
                    relation=EdgeRelation.DERIVES_FROM
                )
                graph.add_edge(edge)
        
        logger.info(
            f"Extracted intent graph with {len(graph.nodes)} nodes "
            f"and {len(graph.edges)} edges"
        )
        
        return graph
    
    async def update_intent(
        self,
        turn: Turn,
        current_graph: IntentGraph
    ) -> IntentGraph:
        """
        Incrementally update intent graph with new turn.
        
        Args:
            turn: New conversation turn
            current_graph: Existing intent graph
            
        Returns:
            Updated IntentGraph
        """
        # Format current graph for prompt
        graph_summary = self._format_graph_summary(current_graph)
        turn_text = f"[{turn.role}]: {turn.content}"
        
        # Call LLM for update
        prompt = INTENT_UPDATE_PROMPT.format(
            intent_graph=graph_summary,
            turn=turn_text
        )
        response = await self.llm.complete(prompt, system=INTENT_EXTRACTION_SYSTEM)
        
        # Parse response
        try:
            data = json.loads(self._clean_json_response(response))
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse intent update response: {e}")
            return current_graph
        
        # Apply status updates
        for update in data.get("status_updates", []):
            node_id = update.get("id")
            new_status = update.get("new_status")
            if node_id in current_graph.nodes and new_status:
                current_graph.nodes[node_id].update_status(normalize_intent_status(new_status))
        
        # Add new nodes
        for node_data in data.get("new_nodes", []):
            # Issue #2 + P4: Robust content extraction from various field names
            content = self._extract_node_content(node_data)
            if not content:
                logger.warning(f"Skipping node without content: {node_data}")
                continue
            
            # Issue #1: Use normalize_intent_type instead of direct IntentType()
            node = IntentNode(
                type=normalize_intent_type(node_data.get("type", "task")),
                content=content,
                status=normalize_intent_status(node_data.get("status", "active")),
                confidence=node_data.get("confidence", 0.8),
                evidence=[turn.id]
            )
            current_graph.add_node(node)
            
            # Handle parent relationship
            parent_id = node_data.get("parent")
            if parent_id and parent_id in current_graph.nodes:
                edge = IntentEdge(
                    source=parent_id,
                    target=node.id,
                    relation=EdgeRelation.DERIVES_FROM
                )
                current_graph.add_edge(edge)
        
        # Add new edges
        for edge_data in data.get("new_edges", []):
            source = edge_data.get("source")
            target = edge_data.get("target")
            relation = edge_data.get("relation", "derives_from")
            
            if source in current_graph.nodes and target in current_graph.nodes:
                edge = IntentEdge(
                    source=source,
                    target=target,
                    relation=normalize_edge_relation(relation)
                )
                current_graph.add_edge(edge)
        
        return current_graph
    
    def get_active_goals(self, graph: IntentGraph) -> list[IntentNode]:
        """Get all currently active goals and tasks."""
        return graph.get_active_goals()
    
    def get_goal_hierarchy(self, graph: IntentGraph, node_id: str) -> list[IntentNode]:
        """Get the full hierarchy path to a specific goal."""
        return graph.get_node_path(node_id)
    
    def _format_conversation(self, turns: list[Turn]) -> str:
        """Format turns for prompt."""
        lines = []
        for i, turn in enumerate(turns):
            lines.append(f"[Turn {i+1}] [{turn.role}]: {turn.content}")
        return "\n\n".join(lines)
    
    def _estimate_tokens(self, text: str) -> int:
        """
        Estimate token count for text.
        
        Uses rough approximation of ~4 characters per token for English text.
        This is conservative to avoid context overflow.
        
        Args:
            text: Text to estimate tokens for
            
        Returns:
            Estimated token count
        """
        return len(text) // 4
    
    async def _summarize_turns(self, turns: list[Turn]) -> str:
        """
        Summarize a batch of conversation turns for intent tracking.
        
        Uses LLM to create a concise summary preserving key goals,
        tasks, and intent changes without all the verbose back-and-forth.
        
        Args:
            turns: List of turns to summarize
            
        Returns:
            Condensed summary string
        """
        if not turns:
            return ""
        
        # Format turns for summarization
        turn_text = []
        for turn in turns:
            turn_text.append(f"[{turn.role}]: {turn.content[:500]}")  # Truncate very long turns
        
        summarization_prompt = f"""Summarize the following conversation segment concisely.
Focus on:
1. Goals mentioned or established
2. Tasks discussed or completed
3. Changes in priorities or direction
4. Key topics and their resolution status

Keep the summary under 500 words. Be factual and specific.

CONVERSATION SEGMENT:
{chr(10).join(turn_text)}

SUMMARY:"""
        
        try:
            summary = await self.llm.complete(summarization_prompt)
            return f"[SUMMARY OF MIDDLE TURNS]: {summary.strip()}"
        except Exception as e:
            logger.warning(f"Failed to summarize turns: {e}")
            # Fallback: just note what was skipped
            return f"[SUMMARY: {len(turns)} turns omitted from middle of conversation]"
    
    async def _format_conversation_smart(
        self,
        turns: list[Turn],
        max_tokens: int = 15000,
        first_n: int = 10,
        last_n: int = 30
    ) -> str:
        """
        Format conversation with smart chunking for long conversations.
        
        Strategy: Sliding Window with Summary
        - Keep first N turns (establishes context, initial goals)
        - Summarize middle turns (preserve key information compactly)
        - Keep last N turns (recent state, current goals)
        
        This handles conversations of any length while preserving
        the most important context for intent extraction.
        
        Args:
            turns: All conversation turns
            max_tokens: Maximum token budget for conversation text
            first_n: Number of initial turns to keep verbatim
            last_n: Number of recent turns to keep verbatim
            
        Returns:
            Formatted conversation string within token budget
        """
        if not turns:
            return ""
        
        total_turns = len(turns)
        
        # For short conversations, use simple formatting
        if total_turns <= (first_n + last_n):
            return self._format_conversation(turns)
        
        # Check if simple formatting fits within budget
        simple_format = self._format_conversation(turns)
        if self._estimate_tokens(simple_format) <= max_tokens:
            return simple_format
        
        logger.info(f"Long conversation detected ({total_turns} turns), applying sliding window")
        
        # Split into three segments
        first_turns = turns[:first_n]
        middle_turns = turns[first_n:-last_n] if last_n > 0 else turns[first_n:]
        last_turns = turns[-last_n:] if last_n > 0 else []
        
        # Format first and last turns verbatim
        first_formatted = []
        for i, turn in enumerate(first_turns):
            first_formatted.append(f"[Turn {i + 1}] [{turn.role}]: {turn.content}")
        
        last_formatted = []
        start_idx = len(turns) - len(last_turns)
        for i, turn in enumerate(last_turns):
            last_formatted.append(f"[Turn {start_idx + i + 1}] [{turn.role}]: {turn.content}")
        
        # Summarize middle section
        middle_summary = await self._summarize_turns(middle_turns)
        
        # Combine all sections
        sections = [
            "=== CONVERSATION START ===",
            "\n\n".join(first_formatted),
            "",
            "=== MIDDLE SECTION (SUMMARIZED) ===",
            middle_summary,
            "",
            "=== RECENT CONVERSATION ===",
            "\n\n".join(last_formatted)
        ]
        
        result = "\n\n".join(sections)
        
        # Log token savings
        original_tokens = self._estimate_tokens(simple_format)
        final_tokens = self._estimate_tokens(result)
        logger.info(f"Conversation chunking: {original_tokens} → {final_tokens} tokens "
                   f"(saved {original_tokens - final_tokens} tokens)")
        
        return result
    
    def _format_graph_summary(self, graph: IntentGraph) -> str:
        """Format graph for update prompt."""
        lines = []
        
        for node_id, node in graph.nodes.items():
            status_emoji = {
                IntentStatus.ACTIVE: "🔵",
                IntentStatus.COMPLETED: "✅",
                IntentStatus.ABANDONED: "❌",
                IntentStatus.PENDING: "⏳",
                IntentStatus.BLOCKED: "🚫"
            }.get(node.status, "")
            
            lines.append(
                f"- [{node.type.value}] {node_id}: {node.content} "
                f"{status_emoji} (confidence: {node.confidence:.2f})"
            )
        
        return "\n".join(lines) if lines else "No existing intents"
    
    def _clean_json_response(self, response: str) -> str:
        """Clean LLM response to extract JSON."""
        # Remove markdown code blocks if present
        response = response.strip()
        if response.startswith("```json"):
            response = response[7:]
        elif response.startswith("```"):
            response = response[3:]
        if response.endswith("```"):
            response = response[:-3]
        return response.strip()

    def _extract_node_content(self, node_data: dict) -> str:
        """
        Extract content from node data with robust field mapping.
        
        P4 Fix: Handle various field names the LLM might use.
        Priority order for content extraction.
        
        Args:
            node_data: Dictionary from LLM response
            
        Returns:
            Extracted content string, or empty string if not found
        """
        # Priority order for content field
        content_fields = [
            "content",      # Standard field
            "description",  # Common alternative
            "label",        # Sometimes used by LLM
            "name",         # Another alternative
            "title",        # Sometimes used for goals
            "text",         # Generic text field
            "summary",      # Occasionally used
            "goal",         # Type-specific
            "task",         # Type-specific
            "decision",     # Type-specific
            "objective",    # Another alternative
        ]
        
        for field in content_fields:
            value = node_data.get(field)
            if value and isinstance(value, str) and value.strip():
                return value.strip()
        
        # Fallback: find the longest string value that's not an ID or status
        excluded_fields = {"id", "type", "status", "parent", "confidence"}
        string_values = [
            (k, v) for k, v in node_data.items()
            if isinstance(v, str) 
            and k not in excluded_fields
            and len(v) > 10  # Skip short values like IDs
        ]
        
        if string_values:
            # Return the longest string
            _, longest = max(string_values, key=lambda x: len(x[1]))
            return longest.strip()
        
        return ""
