"""
Database Engine and Session Management

Provides SQLModel database engine and session handling.
"""

import logging
from contextlib import contextmanager
from pathlib import Path
from typing import Generator

from sqlalchemy import text
from sqlmodel import Session, SQLModel, create_engine

import sys
from shared.config import get_db_connection_string

logger = logging.getLogger(__name__)

# Global engine instance
_engine = None


def get_engine():
    """Get or create the database engine."""
    global _engine
    if _engine is None:
        connection_string = get_db_connection_string()
        # Convert psycopg2 format to SQLAlchemy format
        if connection_string.startswith("postgresql://"):
            # Already in correct format
            pass
        elif "host=" in connection_string:
            # Convert DSN format to URL format
            parts = dict(item.split("=") for item in connection_string.split())
            connection_string = (
                f"postgresql://{parts.get('user', '')}:{parts.get('password', '')}"
                f"@{parts.get('host', 'localhost')}:{parts.get('port', '5432')}"
                f"/{parts.get('dbname', 'docmaster')}"
            )

        _engine = create_engine(
            connection_string,
            echo=False,  # Set to True for SQL debugging
            pool_pre_ping=True,  # Verify connections before use
            pool_size=5,
            max_overflow=10,
        )
    return _engine


def run_migrations() -> None:
    """Run database migrations for new columns."""
    engine = get_engine()

    migrations = [
        # Migration 004: Training datasets tables and dataset_id on training_tasks
        (
            "training_datasets_tables",
            """
            CREATE TABLE IF NOT EXISTS training_datasets (
                dataset_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR(255) NOT NULL,
                description TEXT,
                status VARCHAR(20) NOT NULL DEFAULT 'building',
                train_ratio FLOAT NOT NULL DEFAULT 0.8,
                val_ratio FLOAT NOT NULL DEFAULT 0.1,
                seed INTEGER NOT NULL DEFAULT 42,
                total_documents INTEGER NOT NULL DEFAULT 0,
                total_images INTEGER NOT NULL DEFAULT 0,
                total_annotations INTEGER NOT NULL DEFAULT 0,
                dataset_path VARCHAR(512),
                error_message TEXT,
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
            );
            CREATE INDEX IF NOT EXISTS idx_training_datasets_status ON training_datasets(status);
            """,
        ),
        (
            "dataset_documents_table",
            """
            CREATE TABLE IF NOT EXISTS dataset_documents (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                dataset_id UUID NOT NULL REFERENCES training_datasets(dataset_id) ON DELETE CASCADE,
                document_id UUID NOT NULL REFERENCES admin_documents(document_id),
                split VARCHAR(10) NOT NULL,
                page_count INTEGER NOT NULL DEFAULT 0,
                annotation_count INTEGER NOT NULL DEFAULT 0,
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                UNIQUE(dataset_id, document_id)
            );
            CREATE INDEX IF NOT EXISTS idx_dataset_documents_dataset ON dataset_documents(dataset_id);
            CREATE INDEX IF NOT EXISTS idx_dataset_documents_document ON dataset_documents(document_id);
            """,
        ),
        (
            "training_tasks_dataset_id",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS dataset_id UUID REFERENCES training_datasets(dataset_id);
            CREATE INDEX IF NOT EXISTS idx_training_tasks_dataset ON training_tasks(dataset_id);
            """,
        ),
        # Migration 005: Add group_key to admin_documents
        (
            "admin_documents_group_key",
            """
            ALTER TABLE admin_documents ADD COLUMN IF NOT EXISTS group_key VARCHAR(255);
            CREATE INDEX IF NOT EXISTS ix_admin_documents_group_key ON admin_documents(group_key);
            """,
        ),
        # Migration 006: Model versions table
        (
            "model_versions_table",
            """
            CREATE TABLE IF NOT EXISTS model_versions (
                version_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                version VARCHAR(50) NOT NULL,
                name VARCHAR(255) NOT NULL,
                description TEXT,
                model_path VARCHAR(512) NOT NULL,
                status VARCHAR(20) NOT NULL DEFAULT 'inactive',
                is_active BOOLEAN NOT NULL DEFAULT FALSE,
                task_id UUID REFERENCES training_tasks(task_id),
                dataset_id UUID REFERENCES training_datasets(dataset_id),
                metrics_mAP FLOAT,
                metrics_precision FLOAT,
                metrics_recall FLOAT,
                document_count INTEGER NOT NULL DEFAULT 0,
                training_config JSONB,
                file_size BIGINT,
                trained_at TIMESTAMP WITH TIME ZONE,
                activated_at TIMESTAMP WITH TIME ZONE,
                created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
                updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
            );
            CREATE INDEX IF NOT EXISTS ix_model_versions_version ON model_versions(version);
            CREATE INDEX IF NOT EXISTS ix_model_versions_status ON model_versions(status);
            CREATE INDEX IF NOT EXISTS ix_model_versions_is_active ON model_versions(is_active);
            CREATE INDEX IF NOT EXISTS ix_model_versions_task_id ON model_versions(task_id);
            CREATE INDEX IF NOT EXISTS ix_model_versions_dataset_id ON model_versions(dataset_id);
            """,
        ),
        # Migration 009: Add category to admin_documents
        (
            "admin_documents_category",
            """
            ALTER TABLE admin_documents ADD COLUMN IF NOT EXISTS category VARCHAR(100) DEFAULT 'invoice';
            UPDATE admin_documents SET category = 'invoice' WHERE category IS NULL;
            ALTER TABLE admin_documents ALTER COLUMN category SET NOT NULL;
            CREATE INDEX IF NOT EXISTS idx_admin_documents_category ON admin_documents(category);
            """,
        ),
        # Migration 010: Add training_status and active_training_task_id to training_datasets
        (
            "training_datasets_training_status",
            """
            ALTER TABLE training_datasets ADD COLUMN IF NOT EXISTS training_status VARCHAR(20) DEFAULT NULL;
            ALTER TABLE training_datasets ADD COLUMN IF NOT EXISTS active_training_task_id UUID DEFAULT NULL;
            CREATE INDEX IF NOT EXISTS idx_training_datasets_training_status ON training_datasets(training_status);
            CREATE INDEX IF NOT EXISTS idx_training_datasets_active_training_task_id ON training_datasets(active_training_task_id);
            """,
        ),
        # Migration 010b: Update existing datasets with completed training to 'trained' status
        (
            "training_datasets_update_trained_status",
            """
            UPDATE training_datasets d
            SET status = 'trained'
            WHERE d.status = 'ready'
            AND EXISTS (
                SELECT 1 FROM training_tasks t
                WHERE t.dataset_id = d.dataset_id
                AND t.status = 'completed'
            );
            """,
        ),
        # Migration 007: Add extra columns to training_tasks
        (
            "training_tasks_name",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS name VARCHAR(255);
            UPDATE training_tasks SET name = 'Training ' || substring(task_id::text, 1, 8) WHERE name IS NULL;
            ALTER TABLE training_tasks ALTER COLUMN name SET NOT NULL;
            CREATE INDEX IF NOT EXISTS idx_training_tasks_name ON training_tasks(name);
            """,
        ),
        (
            "training_tasks_description",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS description TEXT;
            """,
        ),
        (
            "training_tasks_admin_token",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS admin_token VARCHAR(255);
            """,
        ),
        (
            "training_tasks_task_type",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS task_type VARCHAR(20) DEFAULT 'train';
            """,
        ),
        (
            "training_tasks_recurring",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS cron_expression VARCHAR(50);
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS is_recurring BOOLEAN DEFAULT FALSE;
            """,
        ),
        (
            "training_tasks_metrics",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS result_metrics JSONB;
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS document_count INTEGER DEFAULT 0;
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS metrics_mAP DOUBLE PRECISION;
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS metrics_precision DOUBLE PRECISION;
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS metrics_recall DOUBLE PRECISION;
            CREATE INDEX IF NOT EXISTS idx_training_tasks_mAP ON training_tasks(metrics_mAP);
            """,
        ),
        (
            "training_tasks_updated_at",
            """
            ALTER TABLE training_tasks ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW();
            """,
        ),
        # Migration 008: Fix model_versions foreign key constraints
        (
            "model_versions_fk_fix",
            """
            ALTER TABLE model_versions DROP CONSTRAINT IF EXISTS model_versions_dataset_id_fkey;
            ALTER TABLE model_versions DROP CONSTRAINT IF EXISTS model_versions_task_id_fkey;
            ALTER TABLE model_versions
                ADD CONSTRAINT model_versions_dataset_id_fkey
                FOREIGN KEY (dataset_id) REFERENCES training_datasets(dataset_id) ON DELETE SET NULL;
            ALTER TABLE model_versions
                ADD CONSTRAINT model_versions_task_id_fkey
                FOREIGN KEY (task_id) REFERENCES training_tasks(task_id) ON DELETE SET NULL;
            """,
        ),
        # Migration 006b: Ensure only one active model at a time
        (
            "model_versions_single_active",
            """
            CREATE UNIQUE INDEX IF NOT EXISTS idx_model_versions_single_active
                ON model_versions(is_active) WHERE is_active = TRUE;
            """,
        ),
    ]

    with engine.connect() as conn:
        for name, sql in migrations:
            try:
                conn.execute(text(sql))
                conn.commit()
                logger.info(f"Migration '{name}' applied successfully")
            except Exception as e:
                # Log but don't fail - column may already exist
                logger.debug(f"Migration '{name}' skipped or failed: {e}")


def create_db_and_tables() -> None:
    """Create all database tables."""
    from backend.data.models import ApiKey, AsyncRequest, RateLimitEvent  # noqa: F401
    from backend.data.admin_models import (  # noqa: F401
        AdminToken,
        AdminDocument,
        AdminAnnotation,
        TrainingTask,
        TrainingLog,
    )

    engine = get_engine()
    SQLModel.metadata.create_all(engine)
    logger.info("Database tables created/verified")

    # Run migrations for new columns
    run_migrations()


def get_session() -> Session:
    """Get a new database session."""
    engine = get_engine()
    return Session(engine)


@contextmanager
def get_session_context() -> Generator[Session, None, None]:
    """Context manager for database sessions with auto-commit/rollback."""
    session = get_session()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def close_engine() -> None:
    """Close the database engine and release connections."""
    global _engine
    if _engine is not None:
        _engine.dispose()
        _engine = None
        logger.info("Database engine closed")


def execute_raw_sql(sql: str) -> None:
    """Execute raw SQL (for migrations)."""
    engine = get_engine()
    with engine.connect() as conn:
        conn.execute(text(sql))
        conn.commit()
