from uuid import UUID
from app.modules.financial.domain.entities import ProjectBudget

class ProjectBudgetRepository:
    """
    Kết nối Entity ProjectBudget xuống CSDL.
    Sử dụng schema 'financial' độc lập.
    """
    def __init__(self, conn):
        self._conn = conn

    def ensure_schema(self):
        with self._conn.cursor() as cur:
            cur.execute("""
                CREATE SCHEMA IF NOT EXISTS financial;
                CREATE TABLE IF NOT EXISTS financial.project_budgets (
                    id UUID PRIMARY KEY,
                    project_id UUID NOT NULL,
                    project_code TEXT NOT NULL,
                    total_allocated NUMERIC(18, 4) DEFAULT 0,
                    status TEXT,
                    created_at TIMESTAMP WITH TIME ZONE
                );
                CREATE INDEX IF NOT EXISTS idx_budget_project_id ON financial.project_budgets(project_id);
            """)

    def save(self, budget: ProjectBudget) -> None:
        self.ensure_schema()
        with self._conn.cursor() as cur:
            cur.execute("""
                INSERT INTO financial.project_budgets 
                (id, project_id, project_code, total_allocated, status, created_at)
                VALUES (%s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET
                    total_allocated = EXCLUDED.total_allocated,
                    status = EXCLUDED.status;
            """, (
                str(budget.id),
                str(budget.project_id),
                budget.project_code,
                budget.total_allocated,
                budget.status,
                budget.created_at
            ))
