"""
Session-scoped DuckDB instances for data cascade temp tables.

Each cascade execution gets its own DuckDB instance where temp tables
persist across cells. Tables are named _<cell_name> by convention.
"""

import os
import duckdb
import atexit
from typing import Dict, Optional, List
from threading import Lock

from ..console_style import S, styled_print


# Global registry of session databases
_session_dbs: Dict[str, duckdb.DuckDBPyConnection] = {}
_session_db_lock = Lock()

# Per-session locks to prevent concurrent access to the same DuckDB connection
# DuckDB connections are NOT thread-safe, so we need to serialize access
_session_locks: Dict[str, Lock] = {}

# Directory for session database files - now in LARS_ROOT
def _get_session_db_dir():
    """Get session database directory (creates if needed)."""
    from ..config import get_config
    config = get_config()
    session_db_dir = os.path.join(config.root_dir, 'session_dbs')
    os.makedirs(session_db_dir, exist_ok=True)
    return session_db_dir


def get_session_db(session_id: str) -> duckdb.DuckDBPyConnection:
    """
    Get or create a DuckDB connection for the given session.

    The database file persists at $LARS_ROOT/session_dbs/<session_id>.duckdb
    and contains all temp tables created during the cascade execution.

    Includes health check: if cached connection is corrupt, it's replaced
    with a fresh one.

    Args:
        session_id: Unique session identifier

    Returns:
        DuckDB connection for this session
    """
    with _session_db_lock:
        # Check if we have a cached connection
        if session_id in _session_dbs:
            conn = _session_dbs[session_id]
            # Health check: verify connection is still usable
            try:
                conn.execute("SELECT 1").fetchone()
                return conn
            except Exception as e:
                # Connection is corrupt/closed - remove and recreate
                print(f"[session_db] [WARN] Cached connection for {session_id} is bad: {e}")
                try:
                    conn.close()
                except:
                    pass
                del _session_dbs[session_id]
                styled_print(f"[session_db] {S.RETRY} Will create fresh connection for {session_id}")

        # Create new connection
        session_db_dir = _get_session_db_dir()

        # Sanitize session_id for filename (replace problematic chars)
        safe_session_id = session_id.replace("/", "_").replace("\\", "_")

        # Create or open session database
        db_path = os.path.join(session_db_dir, f"{safe_session_id}.duckdb")
        conn = duckdb.connect(db_path)

        # Configure for our use case
        conn.execute("SET threads TO 4")

        _session_dbs[session_id] = conn

        # Also create a lock for this session if it doesn't exist
        if session_id not in _session_locks:
            _session_locks[session_id] = Lock()

        return conn


def get_session_lock(session_id: str) -> Lock:
    """
    Get the lock for a session's DuckDB connection.

    IMPORTANT: Use this lock when executing queries on a shared session connection.
    DuckDB connections are NOT thread-safe - concurrent access causes segfaults.

    Usage:
        conn = get_session_db(session_id)
        lock = get_session_lock(session_id)
        with lock:
            result = conn.execute("SELECT ...").fetchdf()

    Args:
        session_id: Session identifier

    Returns:
        Lock for this session
    """
    with _session_db_lock:
        if session_id not in _session_locks:
            _session_locks[session_id] = Lock()
        return _session_locks[session_id]


def cleanup_session_db(session_id: str, delete_file: bool = True):
    """
    Clean up a session's DuckDB resources.

    Called when a cascade completes to free resources.

    Args:
        session_id: Session to clean up
        delete_file: If True, delete the database file
    """
    with _session_db_lock:
        if session_id in _session_dbs:
            conn = _session_dbs.pop(session_id)
            try:
                conn.close()
            except Exception:
                pass

            if delete_file:
                # Sanitize session_id for filename
                safe_session_id = session_id.replace("/", "_").replace("\\", "_")
                session_db_dir = _get_session_db_dir()
                db_path = os.path.join(session_db_dir, f"{safe_session_id}.duckdb")
                try:
                    if os.path.exists(db_path):
                        os.remove(db_path)
                    # Also remove WAL file if present
                    wal_path = db_path + ".wal"
                    if os.path.exists(wal_path):
                        os.remove(wal_path)
                except Exception:
                    pass


def force_close_session(session_id: str):
    """
    Force close a session's DuckDB connection and remove from cache.

    Used when a connection is suspected to be corrupt (e.g., after error during
    disconnect). Does NOT delete the database file - just closes the connection
    so the next access creates a fresh one.

    This is different from cleanup_session_db which is for normal cleanup.

    Args:
        session_id: Session to force close
    """
    with _session_db_lock:
        if session_id in _session_dbs:
            conn = _session_dbs.pop(session_id)
            try:
                # Try to close gracefully
                conn.close()
            except Exception:
                # If close fails, the connection is truly dead - that's fine
                pass
            styled_print(f"[session_db] {S.LINK} Force closed connection for {session_id}")


def list_session_tables(session_id: str) -> List[str]:
    """
    List all tables in a session's DuckDB.

    Useful for debugging and introspection.

    Args:
        session_id: Session to inspect

    Returns:
        List of table names
    """
    try:
        conn = get_session_db(session_id)
        result = conn.execute(
            "SELECT table_name FROM information_schema.tables WHERE table_schema = 'main'"
        ).fetchall()
        return [row[0] for row in result]
    except Exception:
        return []


def get_session_table(session_id: str, table_name: str):
    """
    Get a DataFrame from a session temp table.

    Args:
        session_id: Session ID
        table_name: Table name (with or without _ prefix)

    Returns:
        pandas DataFrame
    """
    conn = get_session_db(session_id)
    # Normalize table name
    if not table_name.startswith('_'):
        table_name = f"_{table_name}"
    return conn.execute(f"SELECT * FROM {table_name}").fetchdf()


def session_db_exists(session_id: str) -> bool:
    """
    Check if a session database exists (in memory or on disk).

    Args:
        session_id: Session to check

    Returns:
        True if session database exists
    """
    with _session_db_lock:
        if session_id in _session_dbs:
            return True

        safe_session_id = session_id.replace("/", "_").replace("\\", "_")
        session_db_dir = _get_session_db_dir()
        db_path = os.path.join(session_db_dir, f"{safe_session_id}.duckdb")
        return os.path.exists(db_path)


# Cleanup on process exit
@atexit.register
def _cleanup_all_sessions():
    """Clean up all session databases on process exit."""
    with _session_db_lock:
        for session_id in list(_session_dbs.keys()):
            try:
                conn = _session_dbs.pop(session_id)
                conn.close()
            except Exception:
                pass
