from __future__ import annotations

from datetime import datetime


class PoolRepositoryMixin:
    def save_reason_override(self, trade_date: str, code: str, reason: str) -> None:
        now = datetime.now().astimezone().isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute(
                """
                INSERT INTO reason_overrides (trade_date, code, reason, updated_at)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(trade_date, code) DO UPDATE SET
                    reason = excluded.reason,
                    updated_at = excluded.updated_at
                """,
                (trade_date, code, reason, now),
            )

    def reason_overrides(self, trade_date: str) -> dict[str, str]:
        with self.connect() as connection:
            rows = connection.execute(
                "SELECT code, reason FROM reason_overrides WHERE trade_date = ?",
                (trade_date,),
            ).fetchall()
        return {row["code"]: row["reason"] for row in rows}