"""
Execution Trace Collection for DSPy Optimization.

Captures agent execution traces (task input, output, success, context)
that serve as training data for DSPy MIPROv2 prompt optimization.

Requirements: SEC-05, SEC-06
"""

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

import json
import sqlite3
from pathlib import Path

# Optional DSPy import with graceful fallback
try:
    import dspy

    HAS_DSPY = True
except ImportError:
    HAS_DSPY = False
    dspy = None


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."""
    import uuid

    return f"trace_{uuid.uuid4().hex[:12]}"


@dataclass
class ExecutionTrace:
    """
    A discrete agent execution trace for optimization training data.

    Captures the full context of an agent's task execution including:
    - What was requested (task_input)
    - What was produced (task_output)
    - Whether it succeeded (success)
    - Contextual information for learning (context)
    - Reward signal for RL-style optimization (reward)

    This forms the foundation for DSPy MIPROv2 prompt optimization.
    """

    # Required identification fields
    trace_id: str
    agent_id: str
    session_id: str

    # Core execution fields
    task_input: Dict[str, Any] = field(default_factory=dict)
    task_output: Dict[str, Any] = field(default_factory=dict)
    success: bool = False
    context: str = ""

    # Optimization fields
    reward: float = 0.0

    # Metadata fields
    timestamp: str = field(default_factory=_utc_now_iso)
    tokens_used: int = 0
    duration_ms: int = 0

    # Optional fields for categorization
    tags: List[str] = field(default_factory=list)

    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,
            "task_input": self.task_input,
            "task_output": self.task_output,
            "success": self.success,
            "context": self.context,
            "reward": self.reward,
            "timestamp": self.timestamp,
            "tokens_used": self.tokens_used,
            "duration_ms": self.duration_ms,
            "tags": self.tags,
        }

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

        Args:
            data: Dictionary with trace fields.

        Returns:
            ExecutionTrace instance.
        """
        return cls(
            trace_id=data.get("trace_id", ""),
            agent_id=data.get("agent_id", ""),
            session_id=data.get("session_id", ""),
            task_input=data.get("task_input", {}),
            task_output=data.get("task_output", {}),
            success=data.get("success", False),
            context=data.get("context", ""),
            reward=data.get("reward", 0.0),
            timestamp=data.get("timestamp", _utc_now_iso()),
            tokens_used=data.get("tokens_used", 0),
            duration_ms=data.get("duration_ms", 0),
            tags=data.get("tags", []),
        )

    def __post_init__(self):
        """Validate and normalize fields after initialization."""
        # Ensure tags is always a list
        if self.tags is None:
            self.tags = []

        # Ensure task_input/output are always dicts
        if self.task_input is None:
            self.task_input = {}
        if self.task_output is None:
            self.task_output = {}

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

    def to_dspy_example(self) -> Optional[Any]:
        """
        Convert trace to DSPy Example for training.

        Returns:
            DSPy Example object if DSPy is available, None otherwise.
        """
        if not HAS_DSPY:
            return None

        # Create a DSPy Example with the trace data
        # The Example is used for MIPROv2 optimization
        example = dspy.Example(
            task_input=json.dumps(self.task_input),
            task_output=json.dumps(self.task_output),
            context=self.context,
            success=str(self.success),
            reward=str(self.reward),
        ).with_inputs("task_input", "context")

        return example

    def add_tag(self, tag: str) -> None:
        """Add a tag to the trace if not already present."""
        if tag not in self.tags:
            self.tags.append(tag)

    def remove_tag(self, tag: str) -> bool:
        """Remove a tag from the trace. Returns True if tag was present."""
        if tag in self.tags:
            self.tags.remove(tag)
            return True
        return False


class TraceCollector:
    """
    SQLite-backed persistent storage for execution traces (SEC-05).

    Provides CRUD operations for traces with indexing on:
    - session_id: for retrieving session history
    - agent_id: for agent-specific trace retrieval
    - success: for filtering successful traces for training
    - timestamp: for time-based queries
    """

    def __init__(self, db_path: str = "optimization/traces.db"):
        """
        Initialize the trace collector.

        Args:
            db_path: Path to SQLite database file. Parent directories
                     will be created if they don't exist.
        """
        self.db_path = Path(db_path)
        self._ensure_tables()

    def _get_connection(self) -> sqlite3.Connection:
        """
        Get a new database connection.

        Returns:
            SQLite connection with row factory enabled.
        """
        return sqlite3.connect(str(self.db_path))

    def _ensure_tables(self) -> None:
        """Initialize database schema with traces table and indexes."""
        # Ensure parent directory exists
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # Create traces table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS traces (
                    trace_id TEXT PRIMARY KEY,
                    agent_id TEXT NOT NULL,
                    session_id TEXT NOT NULL,
                    task_input TEXT DEFAULT '{}',
                    task_output TEXT DEFAULT '{}',
                    success INTEGER DEFAULT 0,
                    context TEXT DEFAULT '',
                    reward REAL DEFAULT 0.0,
                    timestamp TEXT NOT NULL,
                    tokens_used INTEGER DEFAULT 0,
                    duration_ms INTEGER DEFAULT 0,
                    tags TEXT DEFAULT '[]'
                )
            """)

            # Create indexes for efficient queries
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_traces_session_id
                ON traces(session_id)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_traces_agent_id
                ON traces(agent_id)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_traces_success
                ON traces(success)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_traces_timestamp
                ON traces(timestamp)
            """)

            conn.commit()
        finally:
            conn.close()

    def record(self, trace: ExecutionTrace) -> None:
        """
        Store a trace in the database (INSERT OR REPLACE).

        Args:
            trace: ExecutionTrace to persist.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # Serialize task_input, task_output and tags to JSON
            task_input_json = json.dumps(trace.task_input)
            task_output_json = json.dumps(trace.task_output)
            tags_json = json.dumps(trace.tags)

            cursor.execute(
                """
                INSERT OR REPLACE INTO traces (
                    trace_id, agent_id, session_id, task_input, task_output,
                    success, context, reward, timestamp, tokens_used,
                    duration_ms, tags
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    trace.trace_id,
                    trace.agent_id,
                    trace.session_id,
                    task_input_json,
                    task_output_json,
                    1 if trace.success else 0,
                    trace.context,
                    trace.reward,
                    trace.timestamp,
                    trace.tokens_used,
                    trace.duration_ms,
                    tags_json,
                ),
            )

            conn.commit()
        finally:
            conn.close()

    def get(self, trace_id: str) -> Optional[ExecutionTrace]:
        """
        Retrieve a single trace by ID.

        Args:
            trace_id: Unique trace identifier.

        Returns:
            ExecutionTrace if found, None otherwise.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT * FROM traces WHERE trace_id = ?",
                (trace_id,),
            )
            row = cursor.fetchone()
            if row:
                return self._row_to_trace(row)
            return None
        finally:
            conn.close()

    def get_traces_for_agent(
        self, agent_id: str, limit: int = 50
    ) -> List[ExecutionTrace]:
        """
        Get recent traces for an agent.

        Args:
            agent_id: Agent identifier.
            limit: Maximum number of traces to return.

        Returns:
            List of traces for the agent, most recent first.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM traces
                WHERE agent_id = ?
                ORDER BY timestamp DESC
                LIMIT ?
                """,
                (agent_id, limit),
            )
            rows = cursor.fetchall()
            return [self._row_to_trace(row) for row in rows]
        finally:
            conn.close()

    def get_traces_for_session(self, session_id: str) -> List[ExecutionTrace]:
        """
        Get all traces for a session, ordered by timestamp.

        Args:
            session_id: Session identifier.

        Returns:
            List of traces for the session, oldest first.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM traces
                WHERE session_id = ?
                ORDER BY timestamp ASC
                """,
                (session_id,),
            )
            rows = cursor.fetchall()
            return [self._row_to_trace(row) for row in rows]
        finally:
            conn.close()

    def get_successful_traces(self, limit: int = 100) -> List[ExecutionTrace]:
        """
        Get traces that completed successfully.

        These traces are candidates for training data since they
        represent successful task completions.

        Args:
            limit: Maximum number of traces to return.

        Returns:
            List of successful traces, most recent first.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM traces
                WHERE success = 1
                ORDER BY timestamp DESC
                LIMIT ?
                """,
                (limit,),
            )
            rows = cursor.fetchall()
            return [self._row_to_trace(row) for row in rows]
        finally:
            conn.close()

    def get_training_dataset(self, limit: int = 100) -> List[Any]:
        """
        Get successful traces as DSPy Examples for training.

        This method filters for successful traces and converts them
        to DSPy Example objects suitable for MIPROv2 optimization.

        Args:
            limit: Maximum number of traces to return.

        Returns:
            List of DSPy Example objects (empty list if DSPy unavailable).
        """
        if not HAS_DSPY:
            return []

        traces = self.get_successful_traces(limit=limit)
        examples = []
        for trace in traces:
            example = trace.to_dspy_example()
            if example is not None:
                examples.append(example)
        return examples

    def count(self) -> int:
        """
        Get total number of traces in the collector.

        Returns:
            Total trace count.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM traces")
            return cursor.fetchone()[0]
        finally:
            conn.close()

    def delete(self, trace_id: str) -> bool:
        """
        Delete a trace from the collector.

        Args:
            trace_id: Trace to delete.

        Returns:
            True if trace was deleted, False if not found.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                "DELETE FROM traces WHERE trace_id = ?",
                (trace_id,),
            )
            deleted = cursor.rowcount > 0
            conn.commit()
            return deleted
        finally:
            conn.close()

    def _row_to_trace(self, row: tuple) -> ExecutionTrace:
        """
        Convert a database row to an ExecutionTrace dataclass.

        Args:
            row: SQLite row tuple.

        Returns:
            ExecutionTrace instance.
        """
        (
            trace_id,
            agent_id,
            session_id,
            task_input_json,
            task_output_json,
            success,
            context,
            reward,
            timestamp,
            tokens_used,
            duration_ms,
            tags_json,
        ) = row

        # Deserialize JSON fields
        task_input = {}
        if task_input_json:
            try:
                task_input = json.loads(task_input_json)
            except json.JSONDecodeError:
                pass

        task_output = {}
        if task_output_json:
            try:
                task_output = json.loads(task_output_json)
            except json.JSONDecodeError:
                pass

        tags = []
        if tags_json:
            try:
                tags = json.loads(tags_json)
            except json.JSONDecodeError:
                pass

        return ExecutionTrace(
            trace_id=trace_id,
            agent_id=agent_id,
            session_id=session_id,
            task_input=task_input,
            task_output=task_output,
            success=bool(success),
            context=context or "",
            reward=reward or 0.0,
            timestamp=timestamp,
            tokens_used=tokens_used or 0,
            duration_ms=duration_ms or 0,
            tags=tags,
        )
