"""PostgreSQL persistence for Agent command transport."""

from collections.abc import Callable
from typing import Any

import psycopg
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from app.agent_commands.exceptions import AgentCommandConflict, AgentCommandNotFound
from app.config import settings


ConnectionFactory = Callable[[], psycopg.Connection]

CLAIM_COMMANDS_SQL = """
WITH candidates AS (
    SELECT id
    FROM minerops_agent_command_queue
    WHERE agent_id = %s
      AND command_status = 'pending'
    ORDER BY created_at, id
    FOR UPDATE SKIP LOCKED
    LIMIT %s
)
UPDATE minerops_agent_command_queue AS queue
SET command_status = 'claimed', claimed_at = now()
FROM candidates
WHERE queue.id = candidates.id
  AND queue.command_status = 'pending'
RETURNING queue.id AS command_id, queue.command_type,
          queue.request_payload, queue.created_at
"""


def default_connection_factory() -> psycopg.Connection:
    return psycopg.connect(
        host=settings.postgres_host,
        port=settings.postgres_port,
        dbname=settings.postgres_db,
        user=settings.postgres_user,
        password=settings.postgres_password,
        connect_timeout=5,
    )


class AgentCommandRepository:
    """Own queue locking and immutable terminal transitions."""

    def __init__(self, connection_factory: ConnectionFactory | None = None):
        self._connection_factory = connection_factory or default_connection_factory

    def claim_pending(self, agent_id: str, limit: int = 5) -> list[dict]:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(CLAIM_COMMANDS_SQL, (agent_id, limit))
                rows = [dict(row) for row in cur.fetchall()]
            connection.commit()
            return rows
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    def get_command(self, command_id: str) -> dict | None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT id, agent_id, command_type, command_status, request_payload
                    FROM minerops_agent_command_queue WHERE id = %s
                    """,
                    (command_id,),
                )
                row = cur.fetchone()
                return dict(row) if row else None
        finally:
            connection.close()

    def finish_command(
        self,
        command_id: str,
        agent_id: str,
        status: str,
        result_payload: dict[str, Any],
        error_message: str | None,
    ) -> None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    UPDATE minerops_agent_command_queue
                    SET command_status = %s, result_payload = %s,
                        error_message = %s, completed_at = now()
                    WHERE id = %s AND agent_id = %s AND command_status = 'claimed'
                    RETURNING id
                    """,
                    (status, Jsonb(result_payload), error_message, command_id, agent_id),
                )
                if cur.fetchone():
                    connection.commit()
                    return

                cur.execute(
                    """
                    SELECT agent_id, command_status
                    FROM minerops_agent_command_queue WHERE id = %s
                    """,
                    (command_id,),
                )
                row = cur.fetchone()
                connection.rollback()
                if not row:
                    raise AgentCommandNotFound("Agent command not found")
                if row["agent_id"] != agent_id:
                    raise AgentCommandConflict("Agent command belongs to another Agent")
                raise AgentCommandConflict(
                    f"Agent command is {row['command_status']}, not claimed"
                )
        except (AgentCommandNotFound, AgentCommandConflict):
            raise
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()