import os
import logging
import queue
import threading
import time
from typing import Dict, List, Optional, Any
from dotenv import load_dotenv

try:
    import psycopg
    from psycopg import sql as pg_sql
except ImportError:
    psycopg = None
    pg_sql = None

try:
    from psycopg_pool import ConnectionPool
except ImportError:
    ConnectionPool = None

load_dotenv()

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)

_upload_queue: Optional[queue.Queue] = None
_worker_thread: Optional[threading.Thread] = None
_connection_pool: Optional[Any] = None
_connection_pool_lock = threading.Lock()

CHUNK_STATUS_PENDING = "PENDING"


def _get_direct_database_url() -> Optional[str]:
    """Return the direct PostgreSQL connection URL if available."""
    return os.getenv("DIRECT_URL") or os.getenv("DATABASE_URL")


def _get_connection_pool() -> Any:
    db_url = _get_direct_database_url()
    if not db_url:
        raise ValueError("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
    if ConnectionPool is None:
        raise ImportError("psycopg_pool is required for PostgreSQL connection pooling")

    global _connection_pool
    if _connection_pool is None:
        with _connection_pool_lock:
            if _connection_pool is None:
                _connection_pool = ConnectionPool(
                    conninfo=db_url,
                    min_size=1,
                    max_size=20,
                    open=True,
                    kwargs={"connect_timeout": 10},
                )
                logger.info(
                    "Initialized PostgreSQL connection pool",
                    extra={"min_size": 1, "max_size": 20},
                )
    return _connection_pool


def close_connection_pool() -> None:
    """Close the shared PostgreSQL connection pool during application shutdown."""
    global _connection_pool
    with _connection_pool_lock:
        if _connection_pool is None:
            logger.info("No PostgreSQL connection pool to close")
            return

        try:
            _connection_pool.close()
            logger.info("PostgreSQL connection pool closed")
        finally:
            _connection_pool = None


def check_supabase_hash(file_hash: str, table_name: str = "document_chunk") -> bool:
    """
    Check if a document with the given file_hash already exists via PostgreSQL.
    """
    db_url = _get_direct_database_url()
    if not db_url:
        raise ValueError("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
    if psycopg is None:
        raise ImportError("psycopg is required for PostgreSQL connection")

    query = pg_sql.SQL("SELECT 1 FROM {table} WHERE file_hash = %s LIMIT 1").format(
        table=pg_sql.Identifier(table_name)
    )

    try:
        with _get_connection_pool().connection() as conn:
            with conn.cursor() as cur:
                cur.execute(query, (file_hash,))
                exists = cur.fetchone() is not None
                if exists:
                    logger.info(f"Document with hash {file_hash} already exists in DB")
                else:
                    logger.info(f"Document with hash {file_hash} not found in DB")
                return exists
    except Exception as e:
        logger.error(f"Error checking file_hash {file_hash} in table {table_name}: {str(e)}")
        raise


def delete_document_chunks_by_scope(
    file_hash: str,
    user_id: str,
    project_id: str,
    table_names: tuple[str, ...] = ("document_chunk", "document_chunks"),
) -> int:
    """Delete chunk rows for a file within a specific tenant scope."""
    db_url = _get_direct_database_url()
    if not db_url:
        raise ValueError("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
    if psycopg is None:
        raise ImportError("psycopg is required for PostgreSQL connection")

    deleted_count = 0
    try:
        with _get_connection_pool().connection() as conn:
            with conn.cursor() as cur:
                for table_name in table_names:
                    try:
                        query = pg_sql.SQL(
                            "DELETE FROM {table} "
                            "WHERE file_hash = %s AND user_id = %s AND project_id = %s"
                        ).format(table=pg_sql.Identifier(table_name))
                        cur.execute(query, (file_hash, user_id, project_id))
                        deleted_count += cur.rowcount or 0
                    except Exception as exc:
                        logger.warning(
                            "Skipping chunk delete for table",
                            extra={"table_name": table_name, "error": str(exc)},
                        )
        logger.info(
            "Deleted chunk rows by scoped file hash",
            extra={
                "file_hash": file_hash,
                "user_id": user_id,
                "project_id": project_id,
                "deleted_count": deleted_count,
            },
        )
        return deleted_count
    except Exception as e:
        logger.error(
            "Error deleting scoped chunk rows",
            extra={
                "file_hash": file_hash,
                "user_id": user_id,
                "project_id": project_id,
                "error": str(e),
            },
        )
        raise


def insert_chunks_batch(
    chunks: List[Dict[str, Any]], table_name: str = "document_chunk"
) -> Dict[str, Any]:
    """
    Batch insert multiple document chunks into PostgreSQL directly.
    """
    if not chunks:
        logger.warning("No chunks provided for batch insertion")
        return {"inserted": 0, "status": "skipped"}

    db_url = _get_direct_database_url()
    if not db_url:
        raise ValueError("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
    if psycopg is None:
        raise ImportError("psycopg is required for PostgreSQL connection")

    primary_query = pg_sql.SQL(
        """
        INSERT INTO {table} (
            user_id, project_id, file_hash, title, chunk_index, metadata_header, content, embedding_status, error_log
        )
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (user_id, project_id, file_hash, chunk_index) DO NOTHING
        RETURNING id
        """
    ).format(table=pg_sql.Identifier(table_name))

    fallback_query = pg_sql.SQL(
        """
        INSERT INTO {table} (file_hash, title, chunk_index, metadata_header, content)
        VALUES (%s, %s, %s, %s, %s)
        ON CONFLICT (file_hash, chunk_index) DO NOTHING
        RETURNING id
        """
    ).format(table=pg_sql.Identifier(table_name))

    inserted_count = 0
    lifecycle_supported = True
    try:
        with _get_connection_pool().connection() as conn:
            with conn.cursor() as cur:
                for chunk in chunks:
                    if lifecycle_supported:
                        try:
                            cur.execute(
                                primary_query,
                                (
                                    chunk.get("user_id"),
                                    chunk.get("project_id"),
                                    chunk.get("file_hash"),
                                    chunk.get("title"),
                                    chunk.get("chunk_index"),
                                    chunk.get("metadata_header"),
                                    chunk.get("content"),
                                    chunk.get("embedding_status", CHUNK_STATUS_PENDING),
                                    chunk.get("error_log"),
                                ),
                            )
                        except Exception as exc:
                            conn.rollback()
                            lifecycle_supported = False
                            logger.warning(
                                "Chunk lifecycle columns unavailable, falling back to legacy insert",
                                extra={"table_name": table_name, "error": str(exc)},
                            )
                            cur.execute(
                                fallback_query,
                                (
                                    chunk.get("file_hash"),
                                    chunk.get("title"),
                                    chunk.get("chunk_index"),
                                    chunk.get("metadata_header"),
                                    chunk.get("content"),
                                ),
                            )
                    else:
                        cur.execute(
                            fallback_query,
                            (
                                chunk.get("file_hash"),
                                chunk.get("title"),
                                chunk.get("chunk_index"),
                                chunk.get("metadata_header"),
                                chunk.get("content"),
                            ),
                        )

                    if cur.fetchone() is not None:
                        inserted_count += 1
            conn.commit()

        logger.info(f"Successfully inserted {inserted_count} chunks into {table_name}")
        return {
            "inserted": inserted_count,
            "status": "success",
            "response": [{"inserted": inserted_count}],
        }
    except Exception as e:
        error_msg = f"Error inserting {len(chunks)} chunks into {table_name}: {str(e)}"
        logger.error(error_msg)
        raise


def worker_supabase_uploader(
    upload_queue: queue.Queue, table_name: str = "document_chunk", batch_size: int = 10,
    on_batch_complete=None,
) -> None:
    logger.info(f"Starting DB uploader worker (batch_size={batch_size})")

    batch = []
    last_flush_time = time.time()
    flush_timeout = 5

    def _trigger_ghost_worker_wakeup(batch_context: Dict[str, Any] | None = None) -> None:
        """Notify the ghost worker after a successful batch upload.

        The actual wakeup implementation is introduced in the ghost worker
        module in a later phase. This helper is kept defensive so the uploader
        does not fail if the worker-side hook is not yet available.

        Returns:
            None.
        """
        try:
            from app.services.ai.ghost_worker import get_ghost_worker
        except Exception as exc:
            logger.warning(
                "Ghost worker wakeup import unavailable after batch upload",
                extra={"error": str(exc)},
            )
            return

        try:
            ghost_worker = get_ghost_worker()
            trigger_wakeup = getattr(ghost_worker, "trigger_wakeup", None)
            if callable(trigger_wakeup):
                trigger_wakeup()
                logger.info(
                    "Ghost worker wakeup signal dispatched after batch upload",
                    extra={
                        "uploader_thread": threading.current_thread().name,
                        **(batch_context or {}),
                    },
                )
            else:
                logger.warning(
                    "Ghost worker wakeup hook not ready yet",
                    extra={"hook": "trigger_wakeup"},
                )
        except Exception as exc:
            logger.error(
                "Failed to notify GhostWorker after batch upload",
                extra={"error": str(exc)},
            )

    while True:
        queue_item_claimed = False
        try:
            try:
                chunk = upload_queue.get(timeout=1)
                queue_item_claimed = True
            except queue.Empty:
                chunk = None

            if chunk == "STOP":
                if batch:
                    batch_context = {
                        "table_name": table_name,
                        "batch_size": len(batch),
                        "batch_stage": "flush_on_stop",
                    }
                    first = batch[0] if batch else {}
                    batch_context.update(
                        {
                            "file_hash": first.get("file_hash"),
                            "user_id": first.get("user_id"),
                            "project_id": first.get("project_id"),
                        }
                    )
                    logger.info(
                        f"Received STOP signal with {len(batch)} chunks pending. "
                        "Flushing final batch before shutdown..."
                    )
                    try:
                        insert_chunks_batch(batch, table_name)
                        if on_batch_complete:
                            on_batch_complete(batch)
                        _trigger_ghost_worker_wakeup(batch_context)
                    except Exception as e:
                        logger.error(f"Error flushing final batch: {str(e)}")
                logger.info("DB uploader worker shutting down")
                break

            if chunk is not None:
                batch.append(chunk)
                logger.debug(f"Chunk added to batch (current size: {len(batch)})")

            current_time = time.time()
            time_since_flush = current_time - last_flush_time

            if len(batch) >= batch_size or (chunk is None and len(batch) > 0 and time_since_flush > flush_timeout):
                if batch:
                    try:
                        insert_chunks_batch(batch, table_name)
                        logger.info(f"Batch uploaded successfully ({len(batch)} chunks)")
                        if on_batch_complete:
                            on_batch_complete(batch)
                        first = batch[0] if batch else {}
                        _trigger_ghost_worker_wakeup(
                            {
                                "table_name": table_name,
                                "batch_size": len(batch),
                                "batch_stage": "periodic_flush",
                                "file_hash": first.get("file_hash"),
                                "user_id": first.get("user_id"),
                                "project_id": first.get("project_id"),
                            }
                        )
                        batch = []
                        last_flush_time = current_time
                    except Exception as e:
                        logger.error(
                            f"Error uploading batch to DB: {str(e)}. "
                            "Retaining chunks for next attempt."
                        )

        except Exception as e:
            logger.error(f"Unexpected error in uploader worker: {str(e)}")
        finally:
            if queue_item_claimed:
                upload_queue.task_done()


def start_uploader_worker(
    table_name: str = "document_chunk", batch_size: int = 10, on_batch_complete=None,
) -> tuple[queue.Queue, threading.Thread]:
    global _upload_queue, _worker_thread

    if _upload_queue is not None and _worker_thread is not None:
        logger.warning("Uploader worker already running. Returning existing queue and thread.")
        return _upload_queue, _worker_thread

    _upload_queue = queue.Queue()

    _worker_thread = threading.Thread(
        target=worker_supabase_uploader,
        args=(_upload_queue, table_name, batch_size, on_batch_complete),
        daemon=True,
        name="DBUploaderWorker",
    )

    _worker_thread.start()
    logger.info("DB uploader worker thread started")

    return _upload_queue, _worker_thread


def stop_uploader_worker(timeout: int = 10) -> None:
    global _upload_queue, _worker_thread

    if _upload_queue is None or _worker_thread is None:
        logger.warning("No active uploader worker to stop")
        return

    try:
        logger.info("Sending STOP signal to uploader worker...")
        _upload_queue.put("STOP")

        _worker_thread.join(timeout=timeout)

        if _worker_thread.is_alive():
            logger.error(
                f"Worker thread did not stop within {timeout} seconds. "
                "Thread may still be processing."
            )
        else:
            logger.info("Uploader worker thread stopped successfully")
            _upload_queue = None
            _worker_thread = None

    except Exception as e:
        logger.error(f"Error stopping uploader worker: {str(e)}")


def check_document_chunks_exist_by_hash(file_hash: str, table_name: str = "document_chunk") -> bool:
    """
    Check if a document with the given file_hash has any chunks in PostgreSQL.
    Returns True if at least one chunk exists, False otherwise.
    """
    db_url = _get_direct_database_url()
    if not db_url:
        logger.error("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
        return False

    query = pg_sql.SQL("SELECT 1 FROM {table} WHERE file_hash = %s LIMIT 1").format(
        table=pg_sql.Identifier(table_name)
    )

    try:
        with _get_connection_pool().connection() as conn:
            with conn.cursor() as cur:
                cur.execute(query, (file_hash,))
                exists = cur.fetchone() is not None
                return exists
    except Exception as e:
        logger.error(f"Error checking chunks existence for hash {file_hash}: {str(e)}")
        return False


def get_document_chunk_status_summary(
    file_hash: str,
    user_id: str | None = None,
    project_id: str | None = None,
    table_name: str = "document_chunk",
) -> Dict[str, Any]:
    """Return a compact embedding-status summary for one document scope."""
    db_url = _get_direct_database_url()
    if not db_url:
        logger.error("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
        return {
            "file_hash": file_hash,
            "user_id": user_id,
            "project_id": project_id,
            "table_name": table_name,
            "total_chunks": 0,
            "pending_chunks": 0,
            "processing_chunks": 0,
            "completed_chunks": 0,
            "failed_chunks": 0,
        }

    query = pg_sql.SQL(
        """
        SELECT
            COUNT(*) AS total_chunks,
            COUNT(*) FILTER (WHERE embedding_status = 'PENDING') AS pending_chunks,
            COUNT(*) FILTER (WHERE embedding_status = 'PROCESSING') AS processing_chunks,
            COUNT(*) FILTER (WHERE embedding_status = 'COMPLETED') AS completed_chunks,
            COUNT(*) FILTER (WHERE embedding_status = 'FAILED') AS failed_chunks
        FROM {table}
        WHERE file_hash = %s
          AND (%s::text IS NULL OR user_id IS NOT DISTINCT FROM %s)
          AND (%s::text IS NULL OR project_id IS NOT DISTINCT FROM %s)
        """
    ).format(table=pg_sql.Identifier(table_name))

    try:
        with _get_connection_pool().connection() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    query,
                    (file_hash, user_id, user_id, project_id, project_id),
                )
                row = cur.fetchone()
                if row is None:
                    return {
                        "file_hash": file_hash,
                        "user_id": user_id,
                        "project_id": project_id,
                        "table_name": table_name,
                        "total_chunks": 0,
                        "pending_chunks": 0,
                        "processing_chunks": 0,
                        "completed_chunks": 0,
                        "failed_chunks": 0,
                    }

                col_names = [desc.name for desc in cur.description]
                summary = dict(zip(col_names, row))
                summary.update(
                    {
                        "file_hash": file_hash,
                        "user_id": user_id,
                        "project_id": project_id,
                        "table_name": table_name,
                    }
                )
                return summary
    except Exception as e:
        logger.error(
            "Error building document chunk status summary",
            extra={
                "file_hash": file_hash,
                "user_id": user_id,
                "project_id": project_id,
                "table_name": table_name,
                "error": str(e),
            },
        )
        return {
            "file_hash": file_hash,
            "user_id": user_id,
            "project_id": project_id,
            "table_name": table_name,
            "total_chunks": 0,
            "pending_chunks": 0,
            "processing_chunks": 0,
            "completed_chunks": 0,
            "failed_chunks": 0,
        }


def verify_table_exists(table_name: str = "document_chunk") -> bool:
    try:
        if psycopg is None:
            raise ImportError("psycopg is required for PostgreSQL connection")
        db_url = _get_direct_database_url()
        if not db_url:
            raise ValueError("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
        with _get_connection_pool().connection() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT 1
                    FROM information_schema.tables
                    WHERE table_schema = 'public'
                      AND table_name = %s
                    LIMIT 1
                    """,
                    (table_name,),
                )
                exists = cur.fetchone() is not None
                if exists:
                    logger.info(f"Table {table_name} exists and is accessible")
                    return True
                logger.error(f"Table {table_name} not found in PostgreSQL")
                return False
    except Exception as e:
        logger.error(f"Table {table_name} verification failed: {str(e)}")
        return False


async def claim_next_document_batch(
    table_name: str = "document_chunk", limit: int = 20
) -> Optional[Dict[str, Any]]:
    """Claim a bounded batch of chunks for one document scope.

    Args:
        table_name: PostgreSQL table name that stores document chunks.
        limit: Maximum number of chunks to claim in a single batch.

    Returns:
        A dictionary containing the claimed file scope and chunk payloads, or
        None when no eligible batch is available.
    """
    import asyncio
    db_url = _get_direct_database_url()
    if not db_url:
        logger.error("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
        return None

    batch_limit = max(1, limit)

    def _sync_claim():
        try:
            with _get_connection_pool().connection() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        pg_sql.SQL("""
                            WITH anchor AS (
                                SELECT file_hash, user_id, project_id
                                FROM {table}
                                WHERE embedding_status IN ('PENDING', 'FAILED')
                                   OR (embedding_status = 'PROCESSING' AND last_tried_at < NOW() - INTERVAL '5 minutes')
                                ORDER BY COALESCE(last_tried_at, TIMESTAMP 'epoch') ASC, id ASC
                                FOR UPDATE SKIP LOCKED
                                LIMIT 1
                            ),
                            target_batch AS (
                                SELECT c.id
                                FROM {table} c
                                JOIN anchor a ON c.file_hash = a.file_hash
                                             AND c.user_id IS NOT DISTINCT FROM a.user_id
                                             AND c.project_id IS NOT DISTINCT FROM a.project_id
                                WHERE c.embedding_status IN ('PENDING', 'FAILED')
                                   OR (c.embedding_status = 'PROCESSING' AND c.last_tried_at < NOW() - INTERVAL '5 minutes')
                                ORDER BY c.chunk_index ASC, c.id ASC
                                FOR UPDATE SKIP LOCKED
                                LIMIT %s
                            )
                            UPDATE {table}
                            SET embedding_status = 'PROCESSING', last_tried_at = NOW()
                            FROM target_batch
                            WHERE {table}.id = target_batch.id
                            RETURNING 
                                {table}.id, 
                                {table}.file_hash, 
                                {table}.user_id, 
                                {table}.project_id, 
                                {table}.content, 
                                {table}.chunk_index,
                                {table}.title,
                                {table}.metadata_header
                        """).format(table=pg_sql.Identifier(table_name)),
                        (batch_limit,),
                    )
                    rows = cur.fetchall()
                    if not rows:
                        return None
                    
                    # Get column names
                    col_names = [desc.name for desc in cur.description]
                    chunks = [dict(zip(col_names, row)) for row in rows]
                    chunks.sort(
                        key=lambda chunk: (
                            chunk.get("chunk_index") is None,
                            chunk.get("chunk_index"),
                            chunk.get("id"),
                        )
                    )

                    # Return the batch info
                    if chunks:
                        first = chunks[0]
                        logger.info(
                            "Claimed document batch for embedding",
                            extra={
                                "file_hash": first["file_hash"],
                                "user_id": first["user_id"],
                                "project_id": first["project_id"],
                                "claimed_chunks": len(chunks),
                                "table_name": table_name,
                            },
                        )
                        return {
                            "file_hash": first["file_hash"],
                            "user_id": first["user_id"],
                            "project_id": first["project_id"],
                            "chunks": chunks
                        }
                    return None
        except Exception as e:
            logger.error(f"Error claiming document batch: {str(e)}")
            return None

    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, _sync_claim)


async def update_chunks_batch(processed_chunks: List[Dict[str, Any]], table_name: str = "document_chunk") -> bool:
    """
    Bulk update document chunks with their generated embeddings and new status.
    """
    import asyncio
    if not processed_chunks:
        return True

    db_url = _get_direct_database_url()
    if not db_url:
        logger.error("Missing DATABASE_URL or DIRECT_URL for PostgreSQL connection")
        return False

    def _sync_update():
        try:
            with _get_connection_pool().connection() as conn:
                with conn.cursor() as cur:
                    update_query = pg_sql.SQL("""
                        UPDATE {table}
                        SET 
                            embedding = %s,
                            embedding_status = %s,
                            error_log = %s
                        WHERE id = %s
                    """).format(table=pg_sql.Identifier(table_name))
                    
                    data = [
                        (
                            chunk.get("embedding"),
                            chunk.get("embedding_status", "FAILED"),
                            chunk.get("error_log"),
                            chunk.get("id")
                        )
                        for chunk in processed_chunks
                    ]
                    
                    cur.executemany(update_query, data)
                status_completed = sum(
                    1 for chunk in processed_chunks if chunk.get("embedding_status") == "COMPLETED"
                )
                status_failed = sum(
                    1 for chunk in processed_chunks if chunk.get("embedding_status") == "FAILED"
                )
                logger.info(
                    "Updated document chunks batch",
                    extra={
                        "table_name": table_name,
                        "updated_chunks": len(processed_chunks),
                        "status_completed": status_completed,
                        "status_failed": status_failed,
                    },
                )
                return True
        except Exception as e:
            logger.error(f"Error updating chunks batch: {str(e)}")
            return False

    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, _sync_update)
