import json
from typing import Any

from pydantic import model_validator
from sinapsis_core.data_containers.model_types.llm_types import LLMMessage

from sinapsis_chat_history.helpers.database_env_var_keys import DatabaseEnvVars
from sinapsis_chat_history.helpers.schemas import (
    CHAT_HISTORY_INDEXED_COLUMNS,
    ChatHistoryColumn,
    ChatHistoryOrder,
    ChatHistoryOrderBy,
)
from sinapsis_chat_history.providers.base import BaseStorageProvider, DatabaseConfig

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


class PostgresDatabaseConfig(DatabaseConfig):
    """Connection settings for PostgreSQL chat history storage.

    Environment variables are used as fallback only for missing `user`,
    `password`, `host`, and `port` values.
    """

    user: str | None = None
    password: str | None = None
    port: int = 5432
    host: str = "localhost"
    db_name: str = "sinapsis_db"
    sslmode: str | None = None

    @model_validator(mode="before")
    @classmethod
    def load_env_defaults(cls, data: dict[str, Any] | None) -> dict[str, Any]:
        """Use environment variables as fallback for missing connection fields.

        Args:
            data (dict[str, Any] | None): Raw provider config values.

        Returns:
            dict[str, Any]: Config with environment fallbacks applied.
        """
        config = dict(data or {})
        env_map = {
            "user": DatabaseEnvVars.DB_USER,
            "password": DatabaseEnvVars.DB_PASSWORD,
            "host": DatabaseEnvVars.DB_HOST,
            "port": DatabaseEnvVars.DB_PORT,
        }
        for field_name, env_var in env_map.items():
            if config.get(field_name) in {None, ""} and env_var.value:
                config[field_name] = int(env_var.value) if field_name == "port" else env_var.value
        return config

    @model_validator(mode="after")
    def validate_credentials(self) -> "PostgresDatabaseConfig":
        """Require credentials after environment fallback is applied.

        Returns:
            PostgresDatabaseConfig: Validated config instance.
        """
        if not self.user:
            raise ValueError(
                "Database user must be provided via 'DB_USER' environment variable or as a direct attribute."
            )
        if not self.password:
            raise ValueError(
                "Database password must be provided via 'DB_PASSWORD' environment variable or as a direct attribute."
            )
        return self


class PostgresStorageProvider(BaseStorageProvider):
    """PostgreSQL-backed chat history provider."""

    config: PostgresDatabaseConfig

    @staticmethod
    def _ensure_dependency() -> None:
        """Raise a clear error when psycopg is not installed."""
        if psycopg is None or sql is None:
            raise ImportError(
                "PostgreSQL chat history support requires 'psycopg'. Install 'sinapsis-chat-history[psycopg]'."
            )

    def _get_connection_string(self, database: str | None = None) -> str:
        """Build a PostgreSQL connection string.

        Args:
            database (str | None): Optional database name override.

        Returns:
            str: PostgreSQL connection string.
        """
        self._ensure_dependency()
        parts = [
            f"host={self.config.host}",
            f"port={self.config.port}",
            f"user={self.config.user}",
            f"password={self.config.password}",
        ]
        target_db = database or self.config.db_name
        if target_db:
            parts.append(f"dbname={target_db}")
        if self.config.sslmode:
            parts.append(f"sslmode={self.config.sslmode}")
        return " ".join(parts)

    def initialize(self) -> None:
        """Create the target PostgreSQL database and table when missing."""
        self._ensure_dependency()
        default_conn_str = self._get_connection_string(database="postgres")
        with psycopg.connect(default_conn_str, autocommit=True) as conn, conn.cursor() as cur:
            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (self.config.db_name,))
            if not cur.fetchone():
                cur.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(self.config.db_name)))
        self._create_table()

    def _create_table(self) -> None:
        """Create the PostgreSQL table and indexes used for chat history."""
        create_table_stmt = sql.SQL(
            """
            CREATE TABLE IF NOT EXISTS {table} (
                {id} TEXT PRIMARY KEY,
                {user_id} TEXT,
                {session_id} TEXT,
                {role} TEXT NOT NULL,
                {content} TEXT NOT NULL,
                {metadata} JSONB,
                {created_at} TIMESTAMPTZ NOT NULL DEFAULT NOW()
            )
            """
        ).format(
            table=sql.Identifier(self.config.table),
            id=sql.Identifier(ChatHistoryColumn.ID.value),
            user_id=sql.Identifier(ChatHistoryColumn.USER_ID.value),
            session_id=sql.Identifier(ChatHistoryColumn.SESSION_ID.value),
            role=sql.Identifier(ChatHistoryColumn.ROLE.value),
            content=sql.Identifier(ChatHistoryColumn.CONTENT.value),
            metadata=sql.Identifier(ChatHistoryColumn.METADATA.value),
            created_at=sql.Identifier(ChatHistoryColumn.CREATED_AT.value),
        )

        index_statements = [
            sql.SQL("CREATE INDEX IF NOT EXISTS {index_name} ON {table} ({column})").format(
                index_name=sql.Identifier(f"{self.config.table}_{column.value}_idx"),
                table=sql.Identifier(self.config.table),
                column=sql.Identifier(column.value),
            )
            for column in CHAT_HISTORY_INDEXED_COLUMNS
        ]

        with psycopg.connect(self._get_connection_string()) as conn, conn.cursor() as cur:
            cur.execute(create_table_stmt)
            for stmt in index_statements:
                cur.execute(stmt)
            conn.commit()

    @staticmethod
    def _build_where_clause(filters: dict[str, Any]) -> tuple[Any, list[Any]]:
        """Build a parameterized WHERE clause from chat history filters.

        Args:
            filters (dict[str, Any]): Validated filter dictionary.

        Returns:
            tuple[Any, list[Any]]: SQL WHERE clause fragment and bound values.
        """
        clauses: list[Any] = []
        params: list[Any] = []
        for key, value in filters.items():
            clauses.append(sql.SQL("{} = %s").format(sql.Identifier(key)))
            params.append(value)
        if not clauses:
            return sql.SQL("TRUE"), params
        return sql.SQL(" AND ").join(clauses), params

    def save_messages(
        self,
        messages: list[dict[str, Any]],
        metadata: dict[str, Any],
        user_id: str | None = None,
        session_id: str | None = None,
    ) -> int:
        """Persist chat messages to PostgreSQL.

        Args:
            messages (list[dict[str, Any]]): Messages to persist.
            metadata (dict[str, Any]): Shared metadata merged into every stored
                message.
            user_id (str | None): Optional user scope.
            session_id (str | None): Optional session scope.

        Returns:
            int: Number of inserted rows.
        """
        rows = self._normalize_messages(messages=messages, metadata=metadata, user_id=user_id, session_id=session_id)
        insert_stmt = sql.SQL(
            """
            INSERT INTO {table} ({id}, {user_id}, {session_id}, {role}, {content}, {metadata}, {created_at})
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            """
        ).format(
            table=sql.Identifier(self.config.table),
            id=sql.Identifier(ChatHistoryColumn.ID.value),
            user_id=sql.Identifier(ChatHistoryColumn.USER_ID.value),
            session_id=sql.Identifier(ChatHistoryColumn.SESSION_ID.value),
            role=sql.Identifier(ChatHistoryColumn.ROLE.value),
            content=sql.Identifier(ChatHistoryColumn.CONTENT.value),
            metadata=sql.Identifier(ChatHistoryColumn.METADATA.value),
            created_at=sql.Identifier(ChatHistoryColumn.CREATED_AT.value),
        )

        with psycopg.connect(self._get_connection_string()) as conn, conn.cursor() as cur:
            for row in rows:
                cur.execute(
                    insert_stmt,
                    (
                        row[ChatHistoryColumn.ID.value],
                        user_id,
                        session_id,
                        row[ChatHistoryColumn.ROLE.value],
                        row[ChatHistoryColumn.CONTENT.value],
                        json.dumps(row[ChatHistoryColumn.METADATA.value])
                        if row[ChatHistoryColumn.METADATA.value]
                        else None,
                        row[ChatHistoryColumn.CREATED_AT.value],
                    ),
                )
            conn.commit()
        return len(rows)

    def fetch_messages(
        self,
        limit: int,
        offset: int,
        order: ChatHistoryOrder,
        order_by: ChatHistoryOrderBy,
        filters: dict[str, Any],
        user_id: str | None = None,
        session_id: str | None = None,
    ) -> list[LLMMessage]:
        """Load chat messages from PostgreSQL.

        Args:
            limit (int): Maximum number of messages to return.
            offset (int): Number of messages to skip.
            order (ChatHistoryOrder): Sort direction.
            order_by (ChatHistoryOrderBy): Field used for ordering.
            filters (dict[str, Any]): Additional message-level filters.
            user_id (str | None): Optional user scope.
            session_id (str | None): Optional session scope.

        Returns:
            list[LLMMessage]: Normalized chat history messages.
        """
        self._validate_pagination(limit=limit, offset=offset)
        scoped_filters = self._build_scoped_filters(filters=filters, user_id=user_id, session_id=session_id)
        where_clause, params = self._build_where_clause(scoped_filters)
        if limit == -1:
            query = sql.SQL(
                """
                SELECT {role}, {content}, {metadata}
                FROM {table}
                WHERE {where_clause}
                ORDER BY {order_by} {direction}, {id} {direction}
                OFFSET %s
                """
            ).format(
                role=sql.Identifier(ChatHistoryColumn.ROLE.value),
                content=sql.Identifier(ChatHistoryColumn.CONTENT.value),
                metadata=sql.Identifier(ChatHistoryColumn.METADATA.value),
                table=sql.Identifier(self.config.table),
                where_clause=where_clause,
                order_by=sql.Identifier(order_by),
                direction=sql.SQL(order.upper()),
                id=sql.Identifier(ChatHistoryColumn.ID.value),
            )
            params.append(offset)
        else:
            query = sql.SQL(
                """
                SELECT {role}, {content}, {metadata}
                FROM {table}
                WHERE {where_clause}
                ORDER BY {order_by} {direction}, {id} {direction}
                LIMIT %s OFFSET %s
                """
            ).format(
                role=sql.Identifier(ChatHistoryColumn.ROLE.value),
                content=sql.Identifier(ChatHistoryColumn.CONTENT.value),
                metadata=sql.Identifier(ChatHistoryColumn.METADATA.value),
                table=sql.Identifier(self.config.table),
                where_clause=where_clause,
                order_by=sql.Identifier(order_by),
                direction=sql.SQL(order.upper()),
                id=sql.Identifier(ChatHistoryColumn.ID.value),
            )
            params.extend([limit, offset])

        with psycopg.connect(self._get_connection_string()) as conn, conn.cursor() as cur:
            cur.execute(query, params)
            rows = cur.fetchall()

        messages_out: list[LLMMessage] = []
        for role, content, metadata_value in rows:
            payload: dict[str, Any] = {"role": role, "content": content}
            if metadata_value:
                payload["metadata"] = metadata_value
            messages_out.append(LLMMessage(**payload))
        return messages_out

    def remove_messages(
        self,
        filters: dict[str, Any],
        user_id: str | None = None,
        session_id: str | None = None,
    ) -> int:
        """Delete matching chat messages from PostgreSQL.

        Args:
            filters (dict[str, Any]): Additional message-level filters.
            user_id (str | None): Optional user scope.
            session_id (str | None): Optional session scope.

        Returns:
            int: Number of deleted rows.
        """
        scoped_filters = self._build_scoped_filters(filters=filters, user_id=user_id, session_id=session_id)
        where_clause, params = self._build_where_clause(scoped_filters)
        statement = sql.SQL("DELETE FROM {table} WHERE {where_clause}").format(
            table=sql.Identifier(self.config.table),
            where_clause=where_clause,
        )

        with psycopg.connect(self._get_connection_string()) as conn, conn.cursor() as cur:
            cur.execute(statement, params)
            removed = cur.rowcount
            conn.commit()
        return removed

    def reset(self) -> None:
        """Drop and recreate the PostgreSQL chat history table."""
        with psycopg.connect(self._get_connection_string()) as conn, conn.cursor() as cur:
            cur.execute(sql.SQL("DROP TABLE IF EXISTS {}").format(sql.Identifier(self.config.table)))
            conn.commit()
        self._create_table()
