"""Import utility for loading JSON analyses into the database."""

import json
from datetime import datetime
from pathlib import Path
from typing import Any

from loguru import logger
from sqlalchemy.orm import Session

from quber.db.connection import get_session
from quber.db.embeddings import EmbeddingService, get_embedding_service
from quber.db.models import Document, ExtractedTable


def import_json_file(
    json_path: Path | str,
    session: Session | None = None,
    embedding_service: EmbeddingService | None = None,
    generate_embeddings: bool = False,
) -> tuple[Document, int]:
    """
    Import a single JSON analysis file into the database.

    Args:
        json_path: Path to the JSON file.
        session: Optional database session. If None, creates a new session.
        embedding_service: Optional embedding service. If None and generate_embeddings=True, creates one.
        generate_embeddings: Whether to generate embeddings during import.

    If a document with the same filename is already in the database, nothing is
    imported. The existing Document comes back with the number of tables it
    already has in the database.

    Returns:
        Tuple of (Document instance, number of tables imported, or the existing
        document's table count when the import was skipped).

    Raises:
        FileNotFoundError: If the JSON file doesn't exist.
        ValueError: If the JSON structure is invalid.
    """
    json_path = Path(json_path)
    if not json_path.exists():
        raise FileNotFoundError(f"JSON file not found: {json_path}")

    logger.info(f"Importing {json_path.name}...")

    with open(json_path) as f:
        data: dict[str, Any] = json.load(f)

    # Parse extraction date
    extraction_date_str = data.get("extraction_date", "")
    try:
        extraction_date = datetime.fromisoformat(extraction_date_str)
    except ValueError:
        logger.warning(f"Invalid extraction_date format: {extraction_date_str}, using current time")
        extraction_date = datetime.now()

    # Create document record
    document = Document(
        filename=data.get("document", json_path.stem),
        extraction_date=extraction_date,
        total_pages=data.get("total_pages"),
        total_tables=data.get("total_tables"),
        executive_summary=data.get("executive_summary"),
        model_provider=data.get("model_provider"),
        model=data.get("model"),
    )

    # Initialize embedding service if needed
    if generate_embeddings and embedding_service is None:
        logger.info("Initializing embedding service...")
        embedding_service = get_embedding_service()

    # Create table records
    tables_data = data.get("tables", [])
    for table_data in tables_data:
        llm_title = table_data.get("llm_title", "")
        llm_description = table_data.get("llm_description", "")

        # Generate embeddings if requested
        title_emb = None
        desc_emb = None
        if generate_embeddings and llm_title and llm_description:
            assert embedding_service is not None
            # Combined embedding for title + description
            combined_emb = embedding_service.embed_table_metadata(llm_title, llm_description)
            # Store same embedding for both fields (can refine later)
            title_emb = combined_emb.tolist()
            desc_emb = combined_emb.tolist()

        table = ExtractedTable(
            table_id=table_data.get("table_id"),
            page_number=table_data.get("page"),
            procedural_title=table_data.get("procedural_title"),
            llm_title=llm_title,
            llm_description=llm_description,
            table_markdown=table_data.get("table_markdown"),
            headers={"headers": table_data.get("headers", [])},
            table_metadata=table_data.get("metadata", {}),
            title_embedding=title_emb,
            description_embedding=desc_emb,
        )
        document.tables.append(table)

    # Track table count before session operations
    table_count = len(tables_data)

    # Save to database
    if session is None:
        # Use context manager when no session provided
        with get_session() as db_session:
            # Check if document already exists
            existing = db_session.query(Document).filter_by(filename=document.filename).first()
            if existing:
                logger.warning(f"Document {document.filename} already exists (id={existing.id}), skipping")
                existing_count = db_session.query(ExtractedTable).filter_by(document_id=existing.id).count()
                return (existing, existing_count)

            db_session.add(document)
            db_session.commit()
            db_session.refresh(document)

            logger.success(
                f"Imported {document.filename}: {table_count} tables from {document.total_pages} pages"
            )
            return (document, table_count)
    else:
        # Use provided session
        try:
            existing = session.query(Document).filter_by(filename=document.filename).first()
            if existing:
                logger.warning(f"Document {document.filename} already exists (id={existing.id}), skipping")
                existing_count = session.query(ExtractedTable).filter_by(document_id=existing.id).count()
                return (existing, existing_count)

            session.add(document)
            session.commit()
            session.refresh(document)

            logger.success(
                f"Imported {document.filename}: {table_count} tables from {document.total_pages} pages"
            )
            return (document, table_count)

        except Exception as e:
            session.rollback()
            logger.error(f"Failed to import {json_path.name}: {e}")
            raise


def import_directory(
    directory: Path | str,
    pattern: str = "*_analysis.json",
    session: Session | None = None,
    generate_embeddings: bool = False,
) -> tuple[list[Document], int]:
    """
    Import all JSON files matching a pattern from a directory.

    Args:
        directory: Path to directory containing JSON files.
        pattern: Glob pattern for matching files (default: "*_analysis.json").
        session: Optional database session. If None, creates a new session.
        generate_embeddings: Whether to generate embeddings during import.

    Returns:
        Tuple of (list of Document instances, total table count). A skipped file
        whose document already exists contributes that document and its existing
        table count.
    """
    directory = Path(directory)
    if not directory.exists():
        raise FileNotFoundError(f"Directory not found: {directory}")

    json_files = list(directory.glob(pattern))
    if not json_files:
        logger.warning(f"No files matching '{pattern}' found in {directory}")
        return ([], 0)

    logger.info(f"Found {len(json_files)} JSON files to import")

    # Initialize embedding service once for all files if needed
    embedding_service = None
    if generate_embeddings:
        logger.info("Initializing embedding service for batch import...")
        embedding_service = get_embedding_service()

    documents = []
    total_tables = 0

    if session is None:
        # Import each file with its own session
        for json_file in json_files:
            try:
                doc, table_count = import_json_file(
                    json_file,
                    session=None,
                    embedding_service=embedding_service,
                    generate_embeddings=generate_embeddings,
                )
                documents.append(doc)
                total_tables += table_count
            except Exception as e:
                logger.error(f"Failed to import {json_file.name}: {e}")
                continue
    else:
        # Use provided session for all imports
        for json_file in json_files:
            try:
                doc, table_count = import_json_file(
                    json_file,
                    session=session,
                    embedding_service=embedding_service,
                    generate_embeddings=generate_embeddings,
                )
                documents.append(doc)
                total_tables += table_count
            except Exception as e:
                logger.error(f"Failed to import {json_file.name}: {e}")
                continue

    logger.success(f"Successfully imported {len(documents)}/{len(json_files)} files")
    return (documents, total_tables)


def get_import_stats(session: Session | None = None) -> dict[str, Any]:
    """
    Get statistics about imported data.

    Args:
        session: Optional database session. If None, creates a new session.

    Returns:
        Dictionary with import statistics.
    """
    if session is None:
        with get_session() as db_session:
            doc_count = db_session.query(Document).count()
            table_count = db_session.query(ExtractedTable).count()

            # Get documents with table counts
            documents = db_session.query(Document).all()

            stats = {
                "total_documents": doc_count,
                "total_tables": table_count,
                "avg_tables_per_doc": round(table_count / doc_count, 2) if doc_count > 0 else 0,
                "documents": [
                    {"filename": doc.filename, "tables": len(doc.tables), "pages": doc.total_pages}
                    for doc in documents
                ],
            }

            return stats
    else:
        doc_count = session.query(Document).count()
        table_count = session.query(ExtractedTable).count()

        # Get documents with table counts
        documents = session.query(Document).all()

        stats = {
            "total_documents": doc_count,
            "total_tables": table_count,
            "avg_tables_per_doc": round(table_count / doc_count, 2) if doc_count > 0 else 0,
            "documents": [
                {"filename": doc.filename, "tables": len(doc.tables), "pages": doc.total_pages}
                for doc in documents
            ],
        }

        return stats