from __future__ import annotations

from dataclasses import dataclass

from sqlalchemy import text

from app.modules.shared.db import get_engine


@dataclass
class RagJobRow:
    index_job_id: str
    rag_session_id: str
    status: str
    indexed_files: int
    failed_files: int
    cache_hit_files: int
    cache_miss_files: int
    error_code: str | None
    error_desc: str | None
    error_module: str | None


class RagJobRepository:
    def create_job(self, index_job_id: str, rag_session_id: str, status: str) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO rag_index_jobs (index_job_id, rag_session_id, status)
                    VALUES (:jid, :sid, :status)
                    """
                ),
                {"jid": index_job_id, "sid": rag_session_id, "status": status},
            )
            conn.commit()

    def update_job(
        self,
        index_job_id: str,
        *,
        status: str,
        indexed_files: int,
        failed_files: int,
        cache_hit_files: int = 0,
        cache_miss_files: int = 0,
        error_code: str | None = None,
        error_desc: str | None = None,
        error_module: str | None = None,
    ) -> None:
        with get_engine().connect() as conn:
            conn.execute(
                text(
                    """
                    UPDATE rag_index_jobs
                    SET status = :status,
                        indexed_files = :indexed,
                        failed_files = :failed,
                        cache_hit_files = :cache_hit_files,
                        cache_miss_files = :cache_miss_files,
                        error_code = :ecode,
                        error_desc = :edesc,
                        error_module = :emodule,
                        updated_at = CURRENT_TIMESTAMP
                    WHERE index_job_id = :jid
                    """
                ),
                {
                    "jid": index_job_id,
                    "status": status,
                    "indexed": indexed_files,
                    "failed": failed_files,
                    "cache_hit_files": cache_hit_files,
                    "cache_miss_files": cache_miss_files,
                    "ecode": error_code,
                    "edesc": error_desc,
                    "emodule": error_module,
                },
            )
            conn.commit()

    def get_job(self, index_job_id: str) -> RagJobRow | None:
        with get_engine().connect() as conn:
            row = conn.execute(
                text(
                    """
                    SELECT index_job_id, rag_session_id, status, indexed_files, failed_files,
                           cache_hit_files, cache_miss_files, error_code, error_desc, error_module
                    FROM rag_index_jobs
                    WHERE index_job_id = :jid
                    """
                ),
                {"jid": index_job_id},
            ).mappings().fetchone()
        return RagJobRow(**dict(row)) if row else None