import json

from sqlalchemy import text

from app.modules.shared.db import get_engine


class ChatRepository:
    def ensure_tables(self) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    CREATE TABLE IF NOT EXISTS dialog_sessions (
                        dialog_session_id VARCHAR(64) PRIMARY KEY,
                        rag_session_id VARCHAR(64) NOT NULL,
                        created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
                    )
                    """
                )
            )
            conn.execute(
                text(
                    """
                    CREATE TABLE IF NOT EXISTS chat_messages (
                        id BIGSERIAL PRIMARY KEY,
                        dialog_session_id VARCHAR(64) NOT NULL,
                        task_id VARCHAR(64),
                        role VARCHAR(16) NOT NULL,
                        content TEXT NOT NULL,
                        payload JSONB,
                        created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
                    )
                    """
                )
            )
            conn.execute(text("ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS task_id VARCHAR(64)"))
            conn.execute(text("ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS payload JSONB"))
            conn.commit()

    def create_dialog(self, dialog_session_id: str, rag_session_id: str) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO dialog_sessions (dialog_session_id, rag_session_id)
                    VALUES (:did, :sid)
                    """
                ),
                {"did": dialog_session_id, "sid": rag_session_id},
            )
            conn.commit()

    def get_dialog(self, dialog_session_id: str) -> dict | None:
        with get_engine().connect() as conn:
            row = conn.execute(
                text(
                    """
                    SELECT dialog_session_id, rag_session_id
                    FROM dialog_sessions
                    WHERE dialog_session_id = :did
                    """
                ),
                {"did": dialog_session_id},
            ).mappings().fetchone()
        return dict(row) if row else None

    def add_message(
        self,
        dialog_session_id: str,
        role: str,
        content: str,
        task_id: str | None = None,
        payload: dict | None = None,
    ) -> None:
        payload_json = json.dumps(payload, ensure_ascii=False) if payload is not None else None
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO chat_messages (dialog_session_id, task_id, role, content, payload)
                    VALUES (:did, :task_id, :role, :content, CAST(:payload AS JSONB))
                    """
                ),
                {
                    "did": dialog_session_id,
                    "task_id": task_id,
                    "role": role,
                    "content": content,
                    "payload": payload_json,
                },
            )
            conn.commit()