import zlib
from typing import Dict, Generator, List

from sqlalchemy import Boolean, DateTime, Integer, LargeBinary, String, Text
from sqlalchemy.orm import class_mapper

from backend.database.base import db
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)

# Fields in the messages table that contain compressed data
_COMPRESSED_FIELDS = ["artifacts", "trace"]

# Map of table names to their SQLAlchemy model classes
_MODEL_REGISTRY = {}


def _decompress_field(compressed_data: bytes, field_name: str, row_id: str) -> str | None:
    """Decompress a zlib-compressed field and return the JSON string."""
    if not compressed_data:
        return None
    try:
        return zlib.decompress(compressed_data).decode("utf-8")
    except Exception as e:
        logger.warning(f"Failed to decompress {field_name} for message {row_id}: {e}")
        return None


def _get_model_for_table(table_name: str):
    """Get the SQLAlchemy model class for a table name, if it exists."""
    if not _MODEL_REGISTRY:
        # Build registry from all models in db.Model
        for mapper in db.Model.registry.mappers:
            model_class = mapper.class_
            if hasattr(model_class, "__tablename__"):
                _MODEL_REGISTRY[model_class.__tablename__] = model_class

    return _MODEL_REGISTRY.get(table_name)


def list_exportable_tables() -> List[str]:
    """
    List all tables available for export (tables with SQLAlchemy models).

    Returns:
        Sorted list of table names that have corresponding ORM models
    """

    # Build registry if not already built
    if not _MODEL_REGISTRY:
        for mapper in db.Model.registry.mappers:
            model_class = mapper.class_
            if hasattr(model_class, "__tablename__"):
                _MODEL_REGISTRY[model_class.__tablename__] = model_class

    return sorted(_MODEL_REGISTRY.keys())


def get_table_schema(table_name: str) -> List[Dict[str, str]]:
    """
    Get schema information for a table using its SQLAlchemy model.

    Args:
        table_name: Table name

    Returns:
        List of column definitions with name and type

    Raises:
        ValueError: If table doesn't have a corresponding SQLAlchemy model
    """
    model = _get_model_for_table(table_name)

    if not model:
        raise ValueError(
            f"No SQLAlchemy model found for table '{table_name}'. All exportable tables should have models."
        )

    mapper = class_mapper(model)
    columns = []

    for column in mapper.columns:
        col_name = column.name
        col_type = column.type

        # Map SQLAlchemy types to DSS types
        if isinstance(col_type, (String, Text)):
            dss_type = "string"
        elif isinstance(col_type, Integer):
            dss_type = "bigint"
        elif isinstance(col_type, Boolean):
            dss_type = "boolean"
        elif isinstance(col_type, LargeBinary):
            dss_type = "string"
        elif isinstance(col_type, DateTime):
            dss_type = "datetimenotz"
        else:
            dss_type = "string"

        columns.append({"name": col_name, "type": dss_type})

    return columns


def export_table_rows(
    db_url: str, workload_folder_path: str, table_name: str, decompress: bool = True
) -> Generator[Dict, None, None]:
    """
    Export all rows from a table using its ORM model.

    Args:
        table_name: Table name
        decompress: Whether to decompress BLOB data (for messages table)

    Yields:
        Dictionary for each row

    Raises:
        ValueError: If table doesn't have a corresponding SQLAlchemy model
    """
    model = _get_model_for_table(table_name)

    if not model:
        raise ValueError(
            f"No SQLAlchemy model found for table '{table_name}'. All exportable tables should have models."
        )
    from flask import Flask

    from backend.database.base import db
    from backend.setup import setup_app

    # Create a dummy Flask app to get the app context
    app = Flask(__name__)
    # Push an application context manually
    with app.app_context():
        # TODO how to support AHLite ? flag in db ?
        setup_app(app, db_url, workload_folder_path)  # This will initialize db with the app

    with app.app_context():
        with db.session.begin():
            is_messages_table = table_name.endswith("messages")

            # Use ORM query - efficient and handles type conversions automatically
            query = db.session.query(model)

            for instance in query:
                # Convert ORM instance to dictionary
                # row_dict = {col.name: getattr(instance, col.name) for col in class_mapper(model).columns}
                row_dict = {}
                for col in class_mapper(model).columns:
                    val = getattr(instance, col.name)
                    if hasattr(val, "value"):
                        val = val.value
                    row_dict[col.name] = val
                # Handle decompression for compressed fields in messages table
                if decompress and is_messages_table:
                    for field in _COMPRESSED_FIELDS:
                        if field in row_dict:
                            row_dict[field] = _decompress_field(row_dict[field], field, row_dict.get("id"))

                # Convert remaining binary data to hex (skip already decompressed fields)
                for key, value in row_dict.items():
                    if isinstance(value, bytes) and not (
                        decompress and is_messages_table and key in _COMPRESSED_FIELDS
                    ):
                        row_dict[key] = value.hex()

                yield row_dict
