"""
Database Configuration
"""
import os
import logging
import threading
from sqlalchemy import create_engine, event, pool, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from typing import Generator

from app.core.config import settings

logger = logging.getLogger(__name__)

# Run PASSIVE checkpoint on every request so each write is flushed from WAL into main .db
_checkpoint_every_n_requests = 1
_request_counter = 0
_request_counter_lock = threading.Lock()

# SQLite configuration with connection pooling
connect_args = {}
pool_size = 20  # Reduced from 100 - SQLite handles fewer concurrent connections better
max_overflow = 10  # Additional connections beyond pool_size
pool_timeout = 30  # Max seconds to wait for a connection from pool (prevents indefinite blocking)
poolclass = pool.QueuePool  # Use QueuePool for connection pooling

if "sqlite" in settings.DATABASE_URL:
    connect_args = {
        "check_same_thread": False,
        "timeout": 60,  # 60 second timeout for lock acquisition (reduced from 300)
    }
    # Use QueuePool with WAL mode for concurrent read/write operations
    # WAL mode allows multiple readers and one writer simultaneously
    # Note: With WAL mode, we use explicit transactions (not autocommit)
    poolclass = pool.QueuePool

# Create engine with connection pooling
engine = create_engine(
    settings.DATABASE_URL, 
    connect_args=connect_args,
    poolclass=poolclass,
    pool_size=pool_size,
    max_overflow=max_overflow,
    pool_timeout=pool_timeout,  # Added: timeout for getting connection from pool
    pool_pre_ping=True,  # Verify connections before using them
    pool_recycle=1800,  # Recycle connections after 30 minutes (reduced from 1 hour)
    echo=False
)

# Enable WAL mode and optimizations for SQLite
# WAL mode enables concurrent read/write operations with connection pooling
if "sqlite" in settings.DATABASE_URL:
    
    @event.listens_for(engine, "connect")
    def set_sqlite_pragma(dbapi_conn, connection_record):
        """Configure SQLite connection settings for WAL mode and connection pooling"""
        cursor = dbapi_conn.cursor()
        
        # Enable WAL mode for concurrent read/write operations
        # WAL allows multiple readers and one writer simultaneously
        cursor.execute("PRAGMA journal_mode=WAL")
        
        # Set busy timeout to 60 seconds (matches connection timeout)
        # This prevents queries from blocking other requests indefinitely
        cursor.execute("PRAGMA busy_timeout=60000")
        
        # NORMAL = good balance; data is in WAL and durable. Checkpoint below flushes to .db.
        cursor.execute("PRAGMA synchronous=NORMAL")
        cursor.execute("PRAGMA cache_size=10000")  # Increase cache for better performance
        cursor.execute("PRAGMA temp_store=MEMORY")  # Use memory for temp tables
        
        # Additional optimizations for connection pooling
        cursor.execute("PRAGMA foreign_keys=ON")  # Enable foreign key constraints
        cursor.execute("PRAGMA wal_autocheckpoint=1000")  # Auto-checkpoint WAL file every 1000 pages
        
        cursor.close()

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()


from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from contextlib import contextmanager, asynccontextmanager

# Async Engine Configuration
# Re-use settings.DATABASE_URL but inject +aiosqlite
# Force relative path for SQLite to work with aiosqlite if needed
if "sqlite" in settings.DATABASE_URL and "+aiosqlite" not in settings.DATABASE_URL:
    ASYNC_DATABASE_URL = settings.DATABASE_URL.replace("sqlite://", "sqlite+aiosqlite://")
else:
    ASYNC_DATABASE_URL = settings.DATABASE_URL

async_engine = create_async_engine(
    ASYNC_DATABASE_URL,
    connect_args=connect_args,
    # Using default pool (NullPool for SQLite + aiosqlite) to avoid asyncio incompatibility
    pool_pre_ping=True,
    pool_recycle=1800,
    echo=False
)

AsyncSessionLocal = async_sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False
)

@contextmanager
def get_db_session():
    """
    Lazy database session context manager (Sync).
    """
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()

@asynccontextmanager
async def get_async_db_session():
    """
    Lazy database session context manager (Async).
    Usage:
        async with get_async_db_session() as db:
            result = await db.execute(select(Model))
    """
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


# ---------------------------------------------------------------------------
# Document DB (file_sys_metadata.db) - same pooling and WAL as askme360_questions
# ---------------------------------------------------------------------------
_document_engines = {}
_document_engines_lock = threading.Lock()
_document_db_initialized: set[str] = set()
_doc_connect_args = {"check_same_thread": False, "timeout": 60}
_doc_pool_size = 20
_doc_max_overflow = 10


def _document_db_path(db_file: str) -> str:
    """Return absolute path for document DB (cache key)."""
    return os.path.abspath(db_file)


def get_document_engine(db_file: str):
    """Get or create a pooled engine for file_sys_metadata.db (or file_sys_metadata_<id>.db). Uses WAL and QueuePool like main DB."""
    path = _document_db_path(db_file)
    with _document_engines_lock:
        if path not in _document_engines:
            url = f"sqlite:///{path}"
            eng = create_engine(
                url,
                connect_args=_doc_connect_args,
                poolclass=pool.QueuePool,
                pool_size=_doc_pool_size,
                max_overflow=_doc_max_overflow,
                pool_timeout=30,
                pool_pre_ping=True,
                pool_recycle=1800,
                echo=False,
            )

            @event.listens_for(eng, "connect")
            def _set_doc_sqlite_pragma(dbapi_conn, connection_record):
                cursor = dbapi_conn.cursor()
                cursor.execute("PRAGMA journal_mode=WAL")
                cursor.execute("PRAGMA busy_timeout=60000")
                cursor.execute("PRAGMA synchronous=NORMAL")
                cursor.execute("PRAGMA cache_size=10000")
                cursor.execute("PRAGMA temp_store=MEMORY")
                cursor.execute("PRAGMA foreign_keys=ON")
                cursor.execute("PRAGMA wal_autocheckpoint=1000")
                cursor.close()

            _document_engines[path] = eng
        return _document_engines[path]


# SQLite indexes for file_sys_metadata*.db — aligned with DocumentDbService query patterns.
# Each entry: (index_name, table, column_expr, required_columns)
_DOCUMENT_DB_INDEX_SPECS = (
    ("ix_files_folder_path", "files", "folder_path", ("folder_path",)),
    ("ix_files_folder_agent", "files", "folder_path, agent_id", ("folder_path", "agent_id")),
    ("ix_files_agent_upload", "files", "agent_id, upload_date DESC", ("agent_id", "upload_date")),
    ("ix_folders_agent_fullpath", "folders", "agent_id, full_path", ("agent_id", "full_path")),
    ("ix_folders_name_agent", "folders", "folder_name, agent_id", ("folder_name", "agent_id")),
    ("ix_extracted_images_file_id", "extracted_images", "file_id", ("file_id",)),
    (
        "ix_extracted_images_file_agent",
        "extracted_images",
        "file_id, agent_id",
        ("file_id", "agent_id"),
    ),
    (
        "ix_extracted_images_file_page",
        "extracted_images",
        "file_id, page_number, id",
        ("file_id", "page_number", "id"),
    ),
)


def ensure_document_db_indexes(db_file: str) -> None:
    """
    Create indexes on document metadata tables (files, folders, extracted_images)
    to avoid full-table scans on folder listings, agent-scoped queries, and image lookups.
    Safe to call repeatedly (CREATE INDEX IF NOT EXISTS). Skips indexes when columns are missing.
    """
    try:
        eng = get_document_engine(db_file)
        col_cache: dict = {}

        def _table_columns(table: str) -> set:
            if table not in col_cache:
                with eng.connect() as conn:
                    rows = conn.execute(text(f"PRAGMA table_info({table})")).fetchall()
                col_cache[table] = {r[1] for r in rows}
            return col_cache[table]

        with eng.begin() as conn:
            for idx_name, table, col_expr, required in _DOCUMENT_DB_INDEX_SPECS:
                if not _table_columns(table).issuperset(required):
                    continue
                conn.execute(
                    text(f"CREATE INDEX IF NOT EXISTS {idx_name} ON {table} ({col_expr})")
                )
        logger.debug("Document DB indexes ensured: %s", _document_db_path(db_file))
    except Exception as e:
        logger.warning(
            "ensure_document_db_indexes: could not create indexes on %s: %s",
            db_file,
            e,
        )


def init_document_database(db_file: str) -> bool:
    """Create document DB file and tables (files, folders, extracted_images) if missing. Uses WAL. Returns True on success."""
    try:
        path = _document_db_path(db_file)
        with _document_engines_lock:
            if path in _document_db_initialized:
                return True
        db_dir = os.path.dirname(path)
        if db_dir and not os.path.exists(db_dir):
            os.makedirs(db_dir, exist_ok=True)
        eng = get_document_engine(db_file)
        with eng.connect() as conn:
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS files (
                    id TEXT PRIMARY KEY,
                    filename TEXT NOT NULL,
                    original_filename TEXT NOT NULL,
                    file_type TEXT,
                    file_path TEXT NOT NULL,
                    folder_path TEXT DEFAULT '',
                    relative_path TEXT DEFAULT '',
                    upload_date TEXT,
                    file_size INTEGER,
                    content_preview TEXT,
                    full_content TEXT,
                    image_path TEXT,
                    thumbnail_path TEXT,
                    agent_id TEXT
                )
            """))
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS folders (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    folder_name TEXT NOT NULL,
                    folder_path TEXT NOT NULL,
                    parent_path TEXT,
                    full_path TEXT,
                    file_count INTEGER DEFAULT 0,
                    created_date TEXT,
                    agent_id TEXT,
                    UNIQUE(folder_path, agent_id)
                )
            """))
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS extracted_images (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    file_id TEXT NOT NULL,
                    image_path TEXT NOT NULL,
                    image_type TEXT,
                    page_number INTEGER,
                    upload_date TEXT,
                    agent_id TEXT,
                    FOREIGN KEY (file_id) REFERENCES files(id) ON DELETE CASCADE
                )
            """))
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS folder_scan_jobs (
                    id TEXT PRIMARY KEY,
                    root_path TEXT NOT NULL,
                    root_folder_name TEXT NOT NULL,
                    dirs_json TEXT NOT NULL,
                    dir_index INTEGER NOT NULL DEFAULT 0,
                    file_offset INTEGER NOT NULL DEFAULT 0,
                    batch_size INTEGER NOT NULL DEFAULT 100,
                    max_file_size INTEGER,
                    status TEXT NOT NULL DEFAULT 'pending',
                    last_error TEXT,
                    stats_json TEXT,
                    created_at TEXT,
                    updated_at TEXT
                )
            """))
            conn.commit()
        ensure_document_db_indexes(db_file)
        with _document_engines_lock:
            _document_db_initialized.add(path)
        logger.info(f"Document database initialized: {path}")
        return True
    except Exception as e:
        logger.error(f"init_document_database failed: {e}", exc_info=True)
        return False


@contextmanager
def get_document_db_session(db_file: str):
    """Sync context manager for document DB session (pooled, WAL). Use for file_sys_metadata.db operations."""
    eng = get_document_engine(db_file)
    SessionLocalDoc = sessionmaker(autocommit=False, autoflush=False, bind=eng)
    db = SessionLocalDoc()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


def _get_sqlite_path() -> str:
    """Return absolute path to SQLite database file, or empty string if not SQLite."""
    if "sqlite" not in settings.DATABASE_URL:
        return ""
    db_path = settings.DATABASE_URL.replace("sqlite:///", "").lstrip("./")
    return os.path.abspath(db_path)


def checkpoint_wal(mode: str = "PASSIVE", use_direct_connection: bool = False) -> bool:
    """
    Flush WAL (Write-Ahead Log) into the main .db file so data is persisted there.
    use_direct_connection=True uses a new SQLite connection (e.g. after engine.dispose()).
    Modes: PASSIVE (best effort, non-blocking), FULL (wait for readers), TRUNCATE (full + truncate WAL).
    Returns True if checkpoint ran (or not needed), False on error.
    """
    if "sqlite" not in settings.DATABASE_URL:
        return True
    try:
        if use_direct_connection:
            db_path = _get_sqlite_path()
            if not db_path or not os.path.exists(db_path):
                return True
            import sqlite3
            conn = sqlite3.connect(db_path, timeout=60)
            try:
                conn.execute(f"PRAGMA wal_checkpoint({mode})")
                conn.commit()
            finally:
                conn.close()
        else:
            with engine.connect() as conn:
                conn.execute(text(f"PRAGMA wal_checkpoint({mode})"))
                conn.commit()
        logger.debug(f"WAL checkpoint ({mode}) completed")
        return True
    except Exception as e:
        logger.warning(f"WAL checkpoint failed: {e}")
        return False


def get_db() -> Generator[Session, None, None]:
    """
    Database session dependency with proper cleanup.
    - Each request: commit() saves data; WAL checkpoint runs every N requests (never fails the API).
    - Ensures connection is always closed after each API response.
    """
    db = SessionLocal()
    try:
        yield db
        # Commit only if no exception occurred
        try:
            db.commit()
        except Exception:
            db.rollback()
            raise
    except Exception:
        # Rollback on any exception
        db.rollback()
        raise
    finally:
        # Always close the connection after response is sent
        try:
            db.close()
        except Exception:
            # Ignore errors during close, but ensure connection is released
            pass
        # Flush WAL into main .db periodically (best-effort; never fails the API)
        if "sqlite" in settings.DATABASE_URL:
            try:
                global _request_counter
                with _request_counter_lock:
                    _request_counter += 1
                    n = _request_counter
                if n % _checkpoint_every_n_requests == 0:
                    checkpoint_wal("PASSIVE")
            except Exception:
                pass


def init_database() -> bool:
    """
    Initialize database: create file if missing, create all tables.
    Returns True if successful, False otherwise.
    """
    print("Initializing SQLite database...")
    
    try:
        db_path = ""
        # Extract SQLite file path from DATABASE_URL
        if "sqlite" in settings.DATABASE_URL:
            # Format: sqlite:///./path/to/database.db
            # Extract the file path after sqlite:///
            db_path = settings.DATABASE_URL.replace("sqlite:///", "")
            
            # Handle relative paths (starting with ./)
            if db_path.startswith("./"):
                db_path = db_path[2:]
            
            # Get absolute path
            db_path = os.path.abspath(db_path)
            db_dir = os.path.dirname(db_path)
            
            print(f"Database file will be created at: {db_path}")
            print(f"Connection pool initialized with {pool_size} connections")
            print(f"Database file path: {db_path}")
            
            logger.info(f"Initializing database at: {db_path}")
            
            # Create directory if it doesn't exist
            if db_dir and not os.path.exists(db_dir):
                logger.info(f"Creating database directory: {db_dir}")
                os.makedirs(db_dir, exist_ok=True)
            
            # Create database file if it doesn't exist (SQLite will create it automatically)
            if not os.path.exists(db_path):
                logger.info(f"Database file does not exist, will be created: {db_path}")
            else:
                logger.info(f"Database file already exists: {db_path}")
        
        # Create all tables
        logger.info("Creating database tables...")
        Base.metadata.create_all(bind=engine)
        logger.info("Database initialization completed successfully")
        
        if db_path:
             print(f"Database initialized successfully at: {db_path}")
             if os.path.exists(db_path):
                print(f"Database file size: {os.path.getsize(db_path)} bytes")

        # Add any new columns to existing tables
        try:
            all_new_col()
        except Exception as e:
            logger.error(f"Error running all_new_col during init_database: {str(e)}", exc_info=True)
        
        # Flush WAL into main .db so data is persisted there
        if db_path:
            checkpoint_wal("FULL")
        
        return True
        
    except Exception as e:
        logger.error(f"Database initialization failed: {str(e)}", exc_info=True)
        return False


def add_column_if_not_exists(table_name: str, column_name: str, column_datatype: str) -> bool:
    """ Add a column to a table if it doesn't already exist."""
    try:
        # Only works with SQLite
        if "sqlite" not in settings.DATABASE_URL:
            logger.warning("add_column_if_not_exists only supports SQLite databases")
            return False
        
        # Validate table and column names to prevent SQL injection
        # Allow only alphanumeric characters and underscores
        if not table_name.replace("_", "").isalnum():
            logger.error(f"Invalid table name: {table_name}")
            return False
        if not column_name.replace("_", "").isalnum():
            logger.error(f"Invalid column name: {column_name}")
            return False
   
        with engine.connect() as conn:
            # Begin a transaction
            trans = conn.begin()
            try:
                # Check if column exists using PRAGMA table_info
                result = conn.execute(
                    text(f"PRAGMA table_info({table_name})")
                )
                columns = result.fetchall()
                
                # Extract column names (column name is at index 1 in PRAGMA table_info result)
                column_names = [col[1] for col in columns]
                
                if column_name in column_names:
                    logger.info(f"Column '{column_name}' already exists in table '{table_name}'. No action taken.")
                    print(f"Column '{column_name}' already exists in table '{table_name}'. No action taken.")
                    trans.commit()
                    return True
                
                # Column doesn't exist, add it
                logger.info(f"Adding column '{column_name}' to table '{table_name}'...")
                conn.execute(
                    text(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_datatype}")
                )
                trans.commit()
                logger.info(f"Column '{column_name}' added to table '{table_name}' successfully.")
                print(f"Column '{column_name}' added to table '{table_name}' successfully.")
                return True
                
            except Exception as e:
                trans.rollback()
                raise e
                
    except Exception as e:
        logger.error(f"Error adding column '{column_name}' to table '{table_name}': {str(e)}", exc_info=True)
        return False


def all_new_col(): 
    add_column_if_not_exists("training", "agent_id", "TEXT")
    add_column_if_not_exists("questions", "api_endpoint", "TEXT")
    add_column_if_not_exists("questions", "api_params", "TEXT")
    add_column_if_not_exists("questions", "api_response_validator", "TEXT")
    add_column_if_not_exists("questions", "sql_query_summaries", "TEXT")
    add_column_if_not_exists('user_question_log', 'api_endpoint', 'TEXT')
    add_column_if_not_exists('user_question_log', 'api_params', 'TEXT')
    add_column_if_not_exists('user_question_log', 'api_response_validator', 'TEXT')
    add_column_if_not_exists('user_question_log', 'input_tokens', 'INTEGER')
    add_column_if_not_exists('user_question_log', 'output_tokens', 'INTEGER')
    add_column_if_not_exists('user_question_log', 'total_tokens', 'INTEGER')
    add_column_if_not_exists('user_question_log', 'sql_query_summaries', 'TEXT')
    add_column_if_not_exists('agent_management', "database", "TEXT")
    add_column_if_not_exists('agent_management','api', 'TEXT')
    add_column_if_not_exists('agent_management','file', 'TEXT')
    add_column_if_not_exists('agent_management','data_source_type', 'TEXT')
    add_column_if_not_exists('agent_management','default_data_source_type', 'TEXT')
    add_column_if_not_exists('agent_management','db_type', 'TEXT')
    add_column_if_not_exists('agent_management','host', 'TEXT')
    add_column_if_not_exists('agent_management','port', 'TEXT')
    add_column_if_not_exists('agent_management','username', 'TEXT')
    add_column_if_not_exists('agent_management','password', 'TEXT')
    add_column_if_not_exists('agent_management','db_name', 'TEXT')
    add_column_if_not_exists('agent_management','connection', 'TEXT')
    add_column_if_not_exists('agent_management','description', 'TEXT')

    add_column_if_not_exists('files','folder_path','TEXT')
    add_column_if_not_exists('files','relative_path','TEXT')
    add_column_if_not_exists('files','full_content', 'TEXT')

    
    add_column_if_not_exists('files', 'agent_id', 'TEXT')
    add_column_if_not_exists('folders', 'agent_id', 'TEXT')
    add_column_if_not_exists('extracted_images', 'agent_id', 'TEXT')
    add_column_if_not_exists("user_question_log", "question_normalized", "TEXT")
    add_column_if_not_exists("scheduler_status", "current_job_id", "TEXT")
    add_column_if_not_exists("file_schedules", "job_id", "TEXT")
    ensure_sqlite_index_user_question_log_cache()
    pass


def ensure_sqlite_index_user_question_log_cache() -> None:
    """
    Composite index for question reuse cache lookups.
    SQLite only (matches add_column_if_not_exists). For PostgreSQL run manually, e.g.:
      ALTER TABLE user_question_log ADD COLUMN IF NOT EXISTS question_normalized TEXT;
      CREATE INDEX IF NOT EXISTS ix_uql_qn_exec_uid
        ON user_question_log (question_normalized, executed, user_id);
    """
    if "sqlite" not in settings.DATABASE_URL:
        return
    idx_name = "ix_uql_qn_exec_uid"
    try:
        with engine.connect() as conn:
            row = conn.execute(
                text(
                    "SELECT 1 FROM sqlite_master WHERE type='index' AND name=:n"
                ),
                {"n": idx_name},
            ).fetchone()
            if row:
                return
        with engine.begin() as conn:
            conn.execute(
                text(
                    f"CREATE INDEX IF NOT EXISTS {idx_name} "
                    "ON user_question_log (question_normalized, executed, user_id)"
                )
            )
    except Exception as e:
        logger.warning(
            "ensure_sqlite_index_user_question_log_cache: could not create index %s: %s",
            idx_name,
            e,
        )