"""Database operations for tracking migration state."""

import sqlite3
import os
import threading
from contextlib import contextmanager
from datetime import datetime
from typing import Optional, List, Dict, Any
from pathlib import Path


# Valid step names for mark_step_completed()
STEP_COLUMNS = {
    'collection_created',
    'permissions_copied',
    'import_submitted',
    'import_completed',
    'connectors_migrated',
    'chats_migrated',
}


class MigrationDatabase:
    """Manages SQLite database for tracking collection migrations."""

    def __init__(self, db_path: str = "migration_tracking.db"):
        """Initialize database connection and create schema if needed."""
        # Validate database path
        db_path_obj = Path(db_path)
        db_dir = db_path_obj.parent if db_path_obj.parent != Path('.') else Path.cwd()

        if not db_dir.exists():
            raise ValueError(f"Database directory does not exist: {db_dir}")

        if not os.access(db_dir, os.W_OK):
            raise ValueError(f"No write permission for database directory: {db_dir}")

        self.db_path = db_path
        self._lock = threading.Lock()
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row  # Enable column access by name
        self._create_schema()

    @contextmanager
    def _cursor(self):
        """Locked cursor for read operations."""
        with self._lock:
            yield self.conn.cursor()

    @contextmanager
    def _write_cursor(self):
        """Locked cursor for write operations — commits on success, rolls back on failure."""
        with self._lock:
            cursor = self.conn.cursor()
            try:
                yield cursor
                self.conn.commit()
            except Exception:
                self.conn.rollback()
                raise

    def _create_schema(self):
        """Create migrations table if it doesn't exist."""
        cursor = self.conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS collection_migrations (
                old_collection_id TEXT PRIMARY KEY,
                old_collection_name TEXT,
                new_collection_id TEXT,
                new_collection_name TEXT,
                old_model TEXT,
                new_model TEXT,
                job_id TEXT,
                job_status TEXT,
                user_id TEXT,
                username TEXT,
                created_at TIMESTAMP,
                completed_at TIMESTAMP,
                error TEXT,
                collection_created BOOLEAN DEFAULT 0,
                permissions_copied BOOLEAN DEFAULT 0,
                import_submitted BOOLEAN DEFAULT 0,
                import_completed BOOLEAN DEFAULT 0,
                connectors_migrated BOOLEAN DEFAULT 0,
                chats_migrated BOOLEAN DEFAULT 0
            )
        """)

        # Create indexes for better query performance
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_job_status
            ON collection_migrations(job_status)
        """)

        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_username
            ON collection_migrations(username)
        """)

        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_username_status
            ON collection_migrations(username, job_status)
        """)

        self.conn.commit()

    def check_migration_exists(self, collection_id: str) -> Optional[Dict[str, Any]]:
        """
        Check if a collection has already been migrated.

        Returns:
            Dictionary with migration details if exists, None otherwise.
        """
        with self._cursor() as cursor:
            cursor.execute(
                "SELECT * FROM collection_migrations WHERE old_collection_id = ?",
                (collection_id,)
            )
            row = cursor.fetchone()
            return dict(row) if row else None

    def insert_migration(self, migration_record: Dict[str, Any]):
        """Insert a new migration record (or replace if exists).

        Only called for new migrations or force-remigrate (both want a clean slate).
        Incremental updates use mark_step_completed() and update_job_status().
        """
        with self._write_cursor() as cursor:
            cursor.execute("""
                REPLACE INTO collection_migrations (
                    old_collection_id, old_collection_name, new_collection_id,
                    new_collection_name, old_model, new_model, job_id,
                    job_status, user_id, username, created_at, completed_at, error,
                    collection_created, permissions_copied,
                    import_submitted, import_completed, connectors_migrated, chats_migrated
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                migration_record.get('old_collection_id'),
                migration_record.get('old_collection_name'),
                migration_record.get('new_collection_id'),
                migration_record.get('new_collection_name'),
                migration_record.get('old_model'),
                migration_record.get('new_model'),
                migration_record.get('job_id'),
                migration_record.get('job_status', 'pending'),
                migration_record.get('user_id'),
                migration_record.get('username'),
                migration_record.get('created_at'),
                migration_record.get('completed_at'),
                migration_record.get('error'),
                migration_record.get('collection_created', 0),
                migration_record.get('permissions_copied', 0),
                migration_record.get('import_submitted', 0),
                migration_record.get('import_completed', 0),
                migration_record.get('connectors_migrated', 0),
                migration_record.get('chats_migrated', 0),
            ))

    def update_job_status(self, collection_id: str, status: str, job_id: Optional[str] = None, error: Optional[str] = None):
        """Update the job status for a migration."""
        with self._write_cursor() as cursor:
            if job_id:
                cursor.execute("""
                    UPDATE collection_migrations
                    SET job_status = ?, job_id = ?, error = ?
                    WHERE old_collection_id = ?
                """, (status, job_id, error, collection_id))
            else:
                cursor.execute("""
                    UPDATE collection_migrations
                    SET job_status = ?, error = ?
                    WHERE old_collection_id = ?
                """, (status, error, collection_id))

    def mark_completed(self, collection_id: str, completed_at: Optional[datetime] = None):
        if completed_at is None:
            completed_at = datetime.now()

        with self._write_cursor() as cursor:
            cursor.execute("""
                UPDATE collection_migrations
                SET job_status = 'completed', completed_at = ?, import_completed = 1
                WHERE old_collection_id = ?
            """, (completed_at, collection_id))

    def mark_step_completed(self, collection_id: str, step: str):
        """Mark a migration step as completed.

        Args:
            collection_id: The old collection ID (primary key).
            step: Step name — one of: collection_created, permissions_copied,
                  import_submitted, import_completed, connectors_migrated, chats_migrated.
        """
        if step not in STEP_COLUMNS:
            raise ValueError(f"Invalid step: {step}. Must be one of: {STEP_COLUMNS}")

        with self._write_cursor() as cursor:
            # Step name is validated above against STEP_COLUMNS, so this is safe
            cursor.execute(f"""
                UPDATE collection_migrations
                SET {step} = 1
                WHERE old_collection_id = ?
            """, (collection_id,))

    def get_pending_jobs(
        self,
        usernames: Optional[List[str]] = None,
        collection_ids: Optional[List[str]] = None,
        include_completed_needing_post_steps: bool = False,
    ) -> List[Dict[str, Any]]:
        """Get jobs that need attention.

        Args:
            usernames: Filter by usernames.
            collection_ids: Filter by old collection IDs.
            include_completed_needing_post_steps: Also return completed jobs where
                connectors or chats haven't been migrated yet.
        """
        with self._cursor() as cursor:
            if include_completed_needing_post_steps:
                query = """
                    SELECT * FROM collection_migrations
                    WHERE job_id IS NOT NULL
                    AND (
                        job_status IN ('pending', 'running', 'submitted')
                        OR (job_status = 'completed' AND (connectors_migrated = 0 OR chats_migrated = 0))
                    )
                """
            else:
                query = """
                    SELECT * FROM collection_migrations
                    WHERE job_status IN ('pending', 'running', 'submitted')
                    AND job_id IS NOT NULL
                """
            params = []

            if usernames:
                placeholders = ','.join('?' * len(usernames))
                query += f" AND username IN ({placeholders})"
                params.extend(usernames)

            if collection_ids:
                placeholders = ','.join('?' * len(collection_ids))
                query += f" AND old_collection_id IN ({placeholders})"
                params.extend(collection_ids)

            cursor.execute(query, params)
            return [dict(row) for row in cursor.fetchall()]

    def close(self):
        with self._lock:
            if self.conn:
                self.conn.close()
                self.conn = None
