"""SQLite-based persistence for executions and checkpoints."""

import sqlite3
import json
from pathlib import Path
from typing import Any, Dict, List, Optional
from datetime import datetime


def _make_serializable(obj: Any) -> Any:
    """Convert any object to JSON-serializable format.

    Handles FastAPI types (UploadFile, File, Form, etc.) and other non-serializable objects.
    Recursively processes nested structures (lists, dicts, tuples).

    Args:
        obj: The object to make serializable

    Returns:
        A JSON-serializable representation of the object
    """
    # Handle None
    if obj is None:
        return None

    # Try to serialize directly
    try:
        json.dumps(obj)
        return obj
    except (TypeError, ValueError):
        pass

    # Check object type
    obj_type = type(obj).__name__
    obj_module = type(obj).__module__ if hasattr(obj, '__module__') else None

    # Handle FastAPI UploadFile
    if obj_type == 'UploadFile' or (obj_module and 'starlette' in obj_module and 'UploadFile' in obj_type):
        return {
            "__type__": "UploadFile",
            "filename": getattr(obj, 'filename', None),
            "content_type": getattr(obj, 'content_type', None),
            "size": getattr(obj, 'size', None) if hasattr(obj, 'size') else None
        }

    # Handle other FastAPI parameter types (File, Form, etc.)
    if obj_module and 'fastapi' in obj_module:
        return {
            "__type__": obj_type,
            "__module__": obj_module,
            "__str__": str(obj)
        }

    # Handle dictionaries recursively
    if isinstance(obj, dict):
        return {key: _make_serializable(value) for key, value in obj.items()}

    # Handle lists recursively
    if isinstance(obj, (list, tuple)):
        return [_make_serializable(item) for item in obj]

    # Handle datetime objects
    if isinstance(obj, datetime):
        return obj.isoformat()

    # For other non-serializable types, return string representation
    return str(obj)


class DatabaseManager:
    """Manages SQLite database for execution and checkpoint persistence."""

    def __init__(self, db_path: str = ".ombra_executions.db"):
        """Initialize database manager.

        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = Path(db_path)
        self._init_database()

    def _init_database(self) -> None:
        """Initialize database schema if not exists."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS executions (
                    execution_id TEXT PRIMARY KEY,
                    workflow_name TEXT NOT NULL,
                    status TEXT NOT NULL,
                    started_at TIMESTAMP,
                    completed_at TIMESTAMP,
                    error TEXT,
                    inputs TEXT,
                    outputs TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_workflow
                ON executions(workflow_name)
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_status
                ON executions(status)
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_created
                ON executions(created_at DESC)
            """)

            conn.execute("""
                CREATE TABLE IF NOT EXISTS step_executions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    execution_id TEXT NOT NULL,
                    step_name TEXT NOT NULL,
                    sequence_order INTEGER NOT NULL,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    checkpoint_id TEXT,
                    FOREIGN KEY (execution_id) REFERENCES executions(execution_id)
                )
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_execution
                ON step_executions(execution_id, sequence_order)
            """)

            conn.execute("""
                CREATE TABLE IF NOT EXISTS checkpoints (
                    checkpoint_id TEXT PRIMARY KEY,
                    execution_id TEXT,
                    workflow_name TEXT NOT NULL,
                    step_name TEXT NOT NULL,
                    timestamp TEXT NOT NULL,
                    inputs TEXT,
                    outputs TEXT,
                    metadata TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (execution_id) REFERENCES executions(execution_id)
                )
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_workflow_step
                ON checkpoints(workflow_name, step_name, timestamp DESC)
            """)

            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_execution_checkpoints
                ON checkpoints(execution_id)
            """)

            # Workflow schemas table for persisting workflow signatures
            conn.execute("""
                CREATE TABLE IF NOT EXISTS workflow_schemas (
                    workflow_name TEXT PRIMARY KEY,
                    description TEXT,
                    parameters_json TEXT NOT NULL,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            conn.commit()

    def save_execution(self, execution_data: Dict[str, Any]) -> None:
        """Save or update execution metadata.

        Args:
            execution_data: Dictionary with execution fields
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT OR REPLACE INTO executions
                (execution_id, workflow_name, status, started_at, completed_at,
                 error, inputs, outputs)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                execution_data['execution_id'],
                execution_data['workflow_name'],
                execution_data['status'],
                execution_data.get('started_at'),
                execution_data.get('completed_at'),
                execution_data.get('error'),
                json.dumps(_make_serializable(execution_data.get('inputs'))),
                json.dumps(_make_serializable(execution_data.get('outputs')))
            ))
            conn.commit()

    def load_execution(self, execution_id: str) -> Optional[Dict[str, Any]]:
        """Load execution by ID.

        Args:
            execution_id: Execution ID to load

        Returns:
            Dictionary with execution data or None if not found
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT * FROM executions WHERE execution_id = ?
            """, (execution_id,))

            row = cursor.fetchone()
            if not row:
                return None

            return {
                'execution_id': row['execution_id'],
                'workflow_name': row['workflow_name'],
                'status': row['status'],
                'started_at': row['started_at'],
                'completed_at': row['completed_at'],
                'error': row['error'],
                'inputs': json.loads(row['inputs']) if row['inputs'] else {},
                'outputs': json.loads(row['outputs']) if row['outputs'] else None
            }

    def load_all_executions(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """Load all executions, optionally limited.

        Args:
            limit: Maximum number of executions to load (most recent first)

        Returns:
            List of execution dictionaries
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row

            query = "SELECT * FROM executions ORDER BY created_at DESC"
            if limit:
                query += f" LIMIT {limit}"

            cursor = conn.execute(query)

            executions = []
            for row in cursor:
                executions.append({
                    'execution_id': row['execution_id'],
                    'workflow_name': row['workflow_name'],
                    'status': row['status'],
                    'started_at': row['started_at'],
                    'completed_at': row['completed_at'],
                    'error': row['error'],
                    'inputs': json.loads(row['inputs']) if row['inputs'] else {},
                    'outputs': json.loads(row['outputs']) if row['outputs'] else None
                })

            return executions

    def load_executions_for_workflow(self, workflow_name: str) -> List[Dict[str, Any]]:
        """Load all executions for a specific workflow.

        Args:
            workflow_name: Name of workflow

        Returns:
            List of execution dictionaries
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT * FROM executions
                WHERE workflow_name = ?
                ORDER BY created_at DESC
            """, (workflow_name,))

            executions = []
            for row in cursor:
                executions.append({
                    'execution_id': row['execution_id'],
                    'workflow_name': row['workflow_name'],
                    'status': row['status'],
                    'started_at': row['started_at'],
                    'completed_at': row['completed_at'],
                    'error': row['error'],
                    'inputs': json.loads(row['inputs']) if row['inputs'] else {},
                    'outputs': json.loads(row['outputs']) if row['outputs'] else None
                })

            return executions

    def save_checkpoint(self, checkpoint_data: Dict[str, Any]) -> None:
        """Save checkpoint data.

        Args:
            checkpoint_data: Dictionary with checkpoint fields
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT OR REPLACE INTO checkpoints
                (checkpoint_id, execution_id, workflow_name, step_name, timestamp,
                 inputs, outputs, metadata)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                checkpoint_data['checkpoint_id'],
                checkpoint_data.get('execution_id'),
                checkpoint_data['workflow_name'],
                checkpoint_data['step_name'],
                checkpoint_data['timestamp'],
                json.dumps(_make_serializable(checkpoint_data.get('inputs'))),
                json.dumps(_make_serializable(checkpoint_data.get('outputs'))),
                json.dumps(_make_serializable(checkpoint_data.get('metadata', {})))
            ))
            conn.commit()

    def load_checkpoint(self, checkpoint_id: str) -> Optional[Dict[str, Any]]:
        """Load checkpoint by ID.

        Args:
            checkpoint_id: Checkpoint ID to load

        Returns:
            Dictionary with checkpoint data or None if not found
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT * FROM checkpoints WHERE checkpoint_id = ?
            """, (checkpoint_id,))

            row = cursor.fetchone()
            if not row:
                return None

            return {
                'checkpoint_id': row['checkpoint_id'],
                'step_name': row['step_name'],
                'timestamp': row['timestamp'],
                'inputs': json.loads(row['inputs']) if row['inputs'] else {},
                'outputs': json.loads(row['outputs']) if row['outputs'] else None,
                'metadata': json.loads(row['metadata']) if row['metadata'] else {}
            }

    def list_checkpoints(self, workflow_name: str, step_name: Optional[str] = None) -> List[Dict[str, Any]]:
        """List checkpoints for a workflow, optionally filtered by step.

        Args:
            workflow_name: Workflow name
            step_name: Optional step name filter

        Returns:
            List of checkpoint dictionaries
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row

            if step_name:
                cursor = conn.execute("""
                    SELECT * FROM checkpoints
                    WHERE workflow_name = ? AND step_name = ?
                    ORDER BY timestamp DESC
                """, (workflow_name, step_name))
            else:
                cursor = conn.execute("""
                    SELECT * FROM checkpoints
                    WHERE workflow_name = ?
                    ORDER BY timestamp DESC
                """, (workflow_name,))

            checkpoints = []
            for row in cursor:
                checkpoints.append({
                    'checkpoint_id': row['checkpoint_id'],
                    'step_name': row['step_name'],
                    'timestamp': row['timestamp'],
                    'inputs': json.loads(row['inputs']) if row['inputs'] else {},
                    'outputs': json.loads(row['outputs']) if row['outputs'] else None,
                    'metadata': json.loads(row['metadata']) if row['metadata'] else {}
                })

            return checkpoints

    def get_latest_checkpoint(self, workflow_name: str, step_name: str) -> Optional[Dict[str, Any]]:
        """Get the most recent checkpoint for a step.

        Args:
            workflow_name: Workflow name
            step_name: Step name

        Returns:
            Dictionary with checkpoint data or None if not found
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT * FROM checkpoints
                WHERE workflow_name = ? AND step_name = ?
                ORDER BY timestamp DESC
                LIMIT 1
            """, (workflow_name, step_name))

            row = cursor.fetchone()
            if not row:
                return None

            return {
                'checkpoint_id': row['checkpoint_id'],
                'step_name': row['step_name'],
                'timestamp': row['timestamp'],
                'inputs': json.loads(row['inputs']) if row['inputs'] else {},
                'outputs': json.loads(row['outputs']) if row['outputs'] else None,
                'metadata': json.loads(row['metadata']) if row['metadata'] else {}
            }

    def save_step_execution(self, execution_id: str, step_name: str,
                           sequence_order: int, checkpoint_id: str) -> None:
        """Save step execution record.

        Args:
            execution_id: Parent execution ID
            step_name: Name of step
            sequence_order: Order in execution sequence
            checkpoint_id: Associated checkpoint ID
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT INTO step_executions
                (execution_id, step_name, sequence_order, checkpoint_id)
                VALUES (?, ?, ?, ?)
            """, (execution_id, step_name, sequence_order, checkpoint_id))
            conn.commit()

    def load_step_executions(self, execution_id: str) -> List[Dict[str, Any]]:
        """Load step executions for an execution.

        Args:
            execution_id: Execution ID

        Returns:
            List of step execution dictionaries
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT * FROM step_executions
                WHERE execution_id = ?
                ORDER BY sequence_order
            """, (execution_id,))

            steps = []
            for row in cursor:
                steps.append({
                    'step_name': row['step_name'],
                    'sequence_order': row['sequence_order'],
                    'timestamp': row['timestamp'],
                    'checkpoint_id': row['checkpoint_id']
                })

            return steps

    def delete_execution(self, execution_id: str) -> bool:
        """Delete an execution and all related data.

        This will cascade delete all step_executions and checkpoints
        associated with this execution.

        Args:
            execution_id: Execution ID to delete

        Returns:
            True if execution was deleted, False if not found
        """
        with sqlite3.connect(self.db_path) as conn:
            # First check if execution exists
            cursor = conn.execute("""
                SELECT execution_id FROM executions WHERE execution_id = ?
            """, (execution_id,))

            if not cursor.fetchone():
                return False

            # Delete related step executions
            conn.execute("""
                DELETE FROM step_executions WHERE execution_id = ?
            """, (execution_id,))

            # Delete related checkpoints (via execution_id if linked)
            # Note: Currently checkpoints are linked by workflow_name, not execution_id
            # So we only delete the execution and step_executions
            # Checkpoints are kept as they might be needed for other executions

            # Delete the execution itself
            conn.execute("""
                DELETE FROM executions WHERE execution_id = ?
            """, (execution_id,))

            conn.commit()
            return True

    def save_workflow_schema(self, workflow_name: str, description: str,
                            parameters: List[Dict[str, Any]]) -> None:
        """Save or update workflow schema.

        Args:
            workflow_name: Name of the workflow
            description: Workflow description
            parameters: List of parameter dicts with name, type, required, default
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT OR REPLACE INTO workflow_schemas
                (workflow_name, description, parameters_json, updated_at)
                VALUES (?, ?, ?, CURRENT_TIMESTAMP)
            """, (workflow_name, description, json.dumps(parameters)))
            conn.commit()

    def load_workflow_schemas(self) -> List[Dict[str, Any]]:
        """Load all workflow schemas from database.

        Returns:
            List of workflow schema dictionaries
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT workflow_name, description, parameters_json, updated_at
                FROM workflow_schemas
                ORDER BY workflow_name
            """)

            schemas = []
            for row in cursor:
                schemas.append({
                    'workflow_name': row['workflow_name'],
                    'description': row['description'],
                    'parameters': json.loads(row['parameters_json']),
                    'updated_at': row['updated_at']
                })

            return schemas

    def delete_workflow_schema(self, workflow_name: str) -> None:
        """Delete a workflow schema and all its associated data.

        This removes:
        - The workflow schema
        - All executions for this workflow
        - All step executions for those executions

        Args:
            workflow_name: Name of the workflow to delete
        """
        with sqlite3.connect(self.db_path) as conn:
            # Delete step executions first (foreign key constraint)
            conn.execute("""
                DELETE FROM step_executions
                WHERE execution_id IN (
                    SELECT execution_id FROM executions WHERE workflow_name = ?
                )
            """, (workflow_name,))

            # Delete executions
            conn.execute("""
                DELETE FROM executions WHERE workflow_name = ?
            """, (workflow_name,))

            # Delete workflow schema
            conn.execute("""
                DELETE FROM workflow_schemas WHERE workflow_name = ?
            """, (workflow_name,))

            conn.commit()

    def get_workflow_schema(self, workflow_name: str) -> Optional[Dict[str, Any]]:
        """Get a specific workflow schema.

        Args:
            workflow_name: Name of the workflow

        Returns:
            Workflow schema dictionary or None if not found
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT workflow_name, description, parameters_json, updated_at
                FROM workflow_schemas
                WHERE workflow_name = ?
            """, (workflow_name,))

            row = cursor.fetchone()
            if not row:
                return None

            return {
                'workflow_name': row['workflow_name'],
                'description': row['description'],
                'parameters': json.loads(row['parameters_json']),
                'updated_at': row['updated_at']
            }
