from __future__ import annotations

from sqlalchemy import text

from app.modules.rag.persistence.schema_migrator import RagSchemaMigrator
from app.modules.rag.persistence.schema_statements import (
    base_table_statements,
    index_statements,
    legacy_drop_statements,
)
from app.modules.shared.db import get_engine


class RagSchemaRepository:
    def __init__(self) -> None:
        self._migrator = RagSchemaMigrator()

    def ensure_tables(self) -> None:
        engine = get_engine()
        with engine.connect() as conn:
            self._run(conn, base_table_statements())
            self._migrator.migrate(conn)
            self._run(conn, index_statements())
            self._run(conn, legacy_drop_statements())
            conn.commit()

    def _run(self, conn, statements: tuple[str, ...]) -> None:
        for statement in statements:
            conn.execute(text(statement))
