from typing import Optional
from loguru import logger

from src.schemas import Message, SQLQueryExtractor
from src.entities import (
    ApparatusModel,
    IncidentModel,
    PersonnelModel,
    IncidentApparatusModel,
    IncidentPersonnelModel,
    IncidentBaseModel,
)
from src.configs import DatabaseConfig
from src.utils import PydanticAgent, PineconeClient
from sqlalchemy import inspect, text


class DataExtractorService:
    def __init__(self):
        self._table_model_map = {
            "apparatus": ApparatusModel,
            "incident": IncidentModel,
            "personnel": PersonnelModel,
            "auv_incidentapparatus": IncidentApparatusModel,
            "auv_incidentpersonnel": IncidentPersonnelModel,
            "auv_incidentbase": IncidentBaseModel,
        }
        self._pydantic_agent = PydanticAgent
        self._pinecone_client = PineconeClient

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        pass

    async def _model_schema_to_text(self, model_class) -> str:
        """Convert SQLAlchemy model schema to LLM-friendly text."""
        mapper = inspect(model_class)
        table_name = model_class.__tablename__

        lines = [f"Table: {table_name}", "Columns:"]

        for column in mapper.columns:
            col_type = str(column.type)
            nullable = "nullable" if column.nullable else "required"
            pk = " [PRIMARY KEY]" if column.primary_key else ""
            fk = (
                f" [FK -> {list(column.foreign_keys)[0].target_fullname}]"
                if column.foreign_keys
                else ""
            )
            lines.append(f"  - {column.name}: {col_type} ({nullable}){pk}{fk}")

        return "\n".join(lines)

    async def _get_table_schema(self, table_name: str) -> str:
        model_class = self._table_model_map[table_name]
        return await self._model_schema_to_text(model_class)

    async def _get_data(self, sql_query: str):
        try:
            async with DatabaseConfig.async_session() as session:
                result = await session.execute(text(sql_query))
                return result.mappings().all()
        except Exception as e:
            raise e

    async def extract(
        self, user_query: str, message_history: Optional[list[Message]] = []
    ):
        logger.info("Extracting data...")
        async with self._pinecone_client() as pinecone_client:
            vector_results = await pinecone_client.query(query=user_query, n_results=5)
            vector_document_metadata = [
                (doc, meta)
                for doc, meta in zip(
                    vector_results["documents"], vector_results["metadatas"]
                )
                if doc is not None and meta is not None
            ]
            table_schemas = []
            processed_tables = set()
            for met in vector_results["metadatas"]:
                if met["table_name"] in processed_tables:
                    continue
                processed_tables.add(met["table_name"])
                table_schemas.append(await self._get_table_schema(met["table_name"]))

        user_input = f"""
        # User Query: 
        {user_query}

        # Table Schemas: 
        {table_schemas}

        # Description
        {vector_document_metadata}
        """

        async with self._pydantic_agent() as pydantic_agent:
            output: SQLQueryExtractor = await pydantic_agent.run(
                user_input=user_input, message_history=message_history
            )

        sql_query = output.sqlite_query
        data = await self._get_data(sql_query)
        return data, sql_query, output