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

import sqlite3
import os
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.conn = sqlite3.connect(db_path)
        self.conn.row_factory = sqlite3.Row  # Enable column access by name
        self._create_schema()

    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()

    @contextmanager
    def transaction(self):
        """
        Context manager for database transactions.
        Commits if successful, rolls back on any exception.
        """
        try:
            yield self.conn
            self.conn.commit()
        except Exception:
            self.conn.rollback()
            raise

    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.
        """
        cursor = self.conn.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().
        """
        cursor = self.conn.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),
        ))
        self.conn.commit()

    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."""
        cursor = self.conn.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))

        self.conn.commit()

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

        cursor = self.conn.cursor()
        cursor.execute("""
            UPDATE collection_migrations
            SET job_status = 'completed', completed_at = ?, import_completed = 1
            WHERE old_collection_id = ?
        """, (completed_at, collection_id))
        self.conn.commit()

    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}")

        cursor = self.conn.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,))
        self.conn.commit()

    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.
        """
        cursor = self.conn.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):
        if self.conn:
            self.conn.close()
