from __future__ import annotations

import json
from collections.abc import Iterator
from contextlib import contextmanager

import pymysql
from pymysql.cursors import DictCursor

from .config import Settings
from .models import FailureReport


class FailureRepository:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings

    @contextmanager
    def connection(self) -> Iterator[pymysql.Connection]:
        connection = pymysql.connect(
            host=self.settings.database_host,
            port=self.settings.database_port,
            user=self.settings.database_user,
            password=self.settings.database_password,
            database=self.settings.database_name,
            cursorclass=DictCursor,
            autocommit=False,
            connect_timeout=10,
            read_timeout=30,
            write_timeout=30,
        )
        try:
            yield connection
        finally:
            connection.close()

    def claim_next(self) -> FailureReport | None:
        with self.connection() as connection, connection.cursor() as cursor:
            try:
                cursor.execute(
                    """
                    UPDATE failure_reports
                    SET status = 'open', updated_at = NOW()
                    WHERE status = 'agent_running'
                      AND updated_at < DATE_SUB(NOW(), INTERVAL %s MINUTE)
                    """,
                    (self.settings.stale_claim_minutes,),
                )
                connection.commit()
                cursor.execute("SET TRANSACTION ISOLATION LEVEL READ COMMITTED")
                cursor.execute("START TRANSACTION")
            except Exception:
                connection.rollback()
                raise
            cursor.execute(
                """
                SELECT * FROM failure_reports
                WHERE status = 'open'
                ORDER BY id
                LIMIT 1
                FOR UPDATE SKIP LOCKED
                """
            )
            row = cursor.fetchone()
            if not row:
                connection.commit()
                return None
            cursor.execute(
                "UPDATE failure_reports SET status = 'agent_running', updated_at = NOW() WHERE id = %s",
                (row["id"],),
            )
            connection.commit()
            return FailureReport.from_row(row)

    def release(self, report_id: int, reason: str) -> None:
        self._finish(report_id, "open", {"last_dry_run": reason})

    def complete(self, report_id: int, metadata: dict) -> None:
        self._finish(report_id, "resolved", metadata)

    def fail(self, report_id: int, error: str, metadata: dict | None = None) -> None:
        payload = {"agent_error": error, **(metadata or {})}
        self._finish(report_id, "agent_failed", payload)

    def _finish(self, report_id: int, status: str, additions: dict) -> None:
        with self.connection() as connection, connection.cursor() as cursor:
            try:
                cursor.execute("SELECT context FROM failure_reports WHERE id = %s", (report_id,))
                row = cursor.fetchone() or {}
                context = row.get("context") or {}
                if isinstance(context, str):
                    context = json.loads(context)
                context.update(additions)
                cursor.execute(
                    """
                    UPDATE failure_reports
                    SET status = %s, context = %s, updated_at = NOW()
                    WHERE id = %s
                    """,
                    (status, json.dumps(context, ensure_ascii=True), report_id),
                )
                connection.commit()
            except Exception as exception:
                connection.rollback()
                raise RuntimeError(
                    f"Could not finalize failure report {report_id} as {status}"
                ) from exception