"""PostgreSQL persistence for secure Agent enrollment."""

from collections.abc import Callable
from uuid import UUID

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

from app.agent_enrollment.exceptions import (
    AgentEnrollmentConflict,
    AgentEnrollmentCredentialCollision,
    AgentEnrollmentNotFound,
    AgentEnrollmentTokenInvalid,
)
from app.config import settings


ConnectionFactory = Callable[[], psycopg.Connection]

TOKEN_HASH_CONSTRAINT = "uq_agent_enrollments_token_hash"
OPEN_ZONE_CONSTRAINT = "uq_agent_enrollments_open_zone"
AGENT_SECRET_HASH_CONSTRAINT = "uq_minerops_site_agents_agent_secret_hash"
AGENT_ZONE_CONSTRAINT = "uq_minerops_site_agents_agent_zone_id"
AGENT_ID_CONSTRAINT = "minerops_site_agents_pkey"


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 AgentEnrollmentRepository:
    def __init__(self, connection_factory: ConnectionFactory | None = None):
        self._connection_factory = connection_factory or default_connection_factory

    def create_enrollment(
        self,
        *,
        token_hash: str,
        site_id: UUID,
        agent_zone_id: UUID,
        display_name: str,
        actor_user_id: UUID,
    ) -> dict:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT zone.id,
                           zone.site_id,
                           zone.name AS agent_zone_name,
                           zone.is_enabled,
                           site.name AS site_name,
                           agent.agent_id AS assigned_agent_id
                    FROM agent_zones AS zone
                    JOIN sites AS site ON site.id = zone.site_id
                    LEFT JOIN minerops_site_agents AS agent
                      ON agent.agent_zone_id = zone.id
                    WHERE zone.id = %s
                      AND zone.site_id = %s
                    FOR UPDATE OF zone
                    """,
                    (agent_zone_id, site_id),
                )
                zone = cur.fetchone()

                if not zone:
                    raise AgentEnrollmentNotFound(
                        "Agent Zone was not found for the selected Site"
                    )

                if zone["is_enabled"] is not True:
                    raise AgentEnrollmentConflict(
                        "Agent Zone is disabled"
                    )

                if zone["assigned_agent_id"]:
                    raise AgentEnrollmentConflict(
                        "Agent Zone already has an assigned Agent"
                    )

                cur.execute(
                    "SELECT clock_timestamp() AS operation_now"
                )
                operation_now = cur.fetchone()["operation_now"]

                cur.execute(
                    """
                    UPDATE agent_enrollments
                    SET revoked_at = %s
                    WHERE agent_zone_id = %s
                      AND used_at IS NULL
                      AND revoked_at IS NULL
                    RETURNING id
                    """,
                    (operation_now, agent_zone_id),
                )
                replaced_ids = [
                    str(row["id"]) for row in cur.fetchall()
                ]

                cur.execute(
                    """
                    INSERT INTO agent_enrollments (
                        token_hash,
                        site_id,
                        agent_zone_id,
                        display_name,
                        created_by_user_id,
                        created_at,
                        expires_at
                    )
                    VALUES (
                        %s,
                        %s,
                        %s,
                        %s,
                        %s,
                        %s,
                        %s + interval '30 minutes'
                    )
                    RETURNING
                        id,
                        site_id,
                        agent_zone_id,
                        display_name,
                        created_at,
                        expires_at
                    """,
                    (
                        token_hash,
                        site_id,
                        agent_zone_id,
                        display_name,
                        actor_user_id,
                        operation_now,
                        operation_now,
                    ),
                )
                enrollment = dict(cur.fetchone())

                enrollment["site_name"] = zone["site_name"]
                enrollment["agent_zone_name"] = zone["agent_zone_name"]

                cur.execute(
                    """
                    INSERT INTO audit_logs (
                        actor_user_id,
                        action,
                        entity_type,
                        entity_id,
                        new_value_json,
                        created_at
                    )
                    VALUES (
                        %s,
                        'agent.enrollment.created',
                        'agent_enrollment',
                        %s,
                        %s,
                        %s
                    )
                    """,
                    (
                        actor_user_id,
                        enrollment["id"],
                        Jsonb(
                            {
                                "site_id": str(site_id),
                                "agent_zone_id": str(agent_zone_id),
                                "display_name": display_name,
                                "expires_at": enrollment[
                                    "expires_at"
                                ].isoformat(),
                                "replaced_enrollment_ids": replaced_ids,
                            }
                        ),
                        operation_now,
                    ),
                )

            connection.commit()
            return enrollment

        except UniqueViolation as exc:
            connection.rollback()
            constraint = exc.diag.constraint_name

            if constraint == TOKEN_HASH_CONSTRAINT:
                raise AgentEnrollmentCredentialCollision(
                    "Generated enrollment token collided"
                ) from exc

            if constraint == OPEN_ZONE_CONSTRAINT:
                raise AgentEnrollmentConflict(
                    "An active enrollment already exists for this Agent Zone"
                ) from exc

            raise

        except Exception:
            connection.rollback()
            raise

        finally:
            connection.close()

    def list_enrollments(
        self,
        *,
        site_id: UUID | None = None,
        agent_zone_id: UUID | None = None,
        limit: int = 100,
    ) -> list[dict]:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        enrollment.id,
                        enrollment.site_id,
                        site.name AS site_name,
                        enrollment.agent_zone_id,
                        zone.name AS agent_zone_name,
                        enrollment.display_name,
                        enrollment.created_by_user_id,
                        enrollment.created_at,
                        enrollment.expires_at,
                        enrollment.used_at,
                        enrollment.revoked_at,
                        enrollment.enrolled_agent_id,
                        CASE
                            WHEN enrollment.used_at IS NOT NULL
                                THEN 'used'
                            WHEN enrollment.revoked_at IS NOT NULL
                                THEN 'revoked'
                            WHEN enrollment.expires_at <= now()
                                THEN 'expired'
                            ELSE 'waiting'
                        END AS status
                    FROM agent_enrollments AS enrollment
                    JOIN sites AS site
                      ON site.id = enrollment.site_id
                    JOIN agent_zones AS zone
                      ON zone.id = enrollment.agent_zone_id
                    WHERE (%s::uuid IS NULL OR enrollment.site_id = %s::uuid)
                      AND (
                          %s::uuid IS NULL
                          OR enrollment.agent_zone_id = %s::uuid
                      )
                    ORDER BY enrollment.created_at DESC, enrollment.id DESC
                    LIMIT %s
                    """,
                    (
                        site_id,
                        site_id,
                        agent_zone_id,
                        agent_zone_id,
                        limit,
                    ),
                )
                return [dict(row) for row in cur.fetchall()]
        finally:
            connection.close()

    def redeem_enrollment(
        self,
        *,
        token_hash: str,
        agent_id: str,
        agent_secret_hash: str,
    ) -> dict:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        id,
                        site_id,
                        agent_zone_id,
                        display_name,
                        expires_at,
                        used_at,
                        revoked_at
                    FROM agent_enrollments
                    WHERE token_hash = %s
                    FOR UPDATE
                    """,
                    (token_hash,),
                )
                enrollment = cur.fetchone()

                if not enrollment:
                    raise AgentEnrollmentTokenInvalid(
                        "Invalid or expired enrollment token"
                    )

                if (
                    enrollment["used_at"] is not None
                    or enrollment["revoked_at"] is not None
                ):
                    raise AgentEnrollmentTokenInvalid(
                        "Invalid or expired enrollment token"
                    )

                cur.execute(
                    """
                    SELECT
                        zone.id,
                        zone.site_id,
                        zone.name AS agent_zone_name,
                        zone.is_enabled,
                        site.name AS site_name
                    FROM agent_zones AS zone
                    JOIN sites AS site ON site.id = zone.site_id
                    WHERE zone.id = %s
                      AND zone.site_id = %s
                    FOR UPDATE OF zone
                    """,
                    (
                        enrollment["agent_zone_id"],
                        enrollment["site_id"],
                    ),
                )
                zone = cur.fetchone()

                if not zone:
                    raise AgentEnrollmentConflict(
                        "Enrollment Agent Zone no longer exists"
                    )

                if zone["is_enabled"] is not True:
                    raise AgentEnrollmentConflict(
                        "Agent Zone is disabled"
                    )

                cur.execute(
                    """
                    SELECT agent_id
                    FROM minerops_site_agents
                    WHERE agent_zone_id = %s
                    FOR UPDATE
                    """,
                    (enrollment["agent_zone_id"],),
                )
                existing_agent = cur.fetchone()

                if existing_agent:
                    raise AgentEnrollmentConflict(
                        "Agent Zone already has an assigned Agent"
                    )

                cur.execute(
                    "SELECT clock_timestamp() AS operation_now"
                )
                operation_now = cur.fetchone()["operation_now"]

                if enrollment["expires_at"] <= operation_now:
                    raise AgentEnrollmentTokenInvalid(
                        "Invalid or expired enrollment token"
                    )

                cur.execute(
                    """
                    INSERT INTO minerops_site_agents (
                        agent_id,
                        site_id,
                        agent_zone_id,
                        site_name,
                        display_name,
                        agent_secret_hash,
                        credential_created_at
                    )
                    VALUES (
                        %s,
                        %s,
                        %s,
                        %s,
                        %s,
                        %s,
                        %s
                    )
                    """,
                    (
                        agent_id,
                        enrollment["site_id"],
                        enrollment["agent_zone_id"],
                        zone["site_name"],
                        enrollment["display_name"],
                        agent_secret_hash,
                        operation_now,
                    ),
                )

                cur.execute(
                    """
                    UPDATE agent_enrollments
                    SET used_at = %s,
                        enrolled_agent_id = %s
                    WHERE id = %s
                    RETURNING used_at
                    """,
                    (operation_now, agent_id, enrollment["id"]),
                )
                used_at = cur.fetchone()["used_at"]

                cur.execute(
                    """
                    SELECT network::text AS network
                    FROM agent_zone_network_ranges
                    WHERE agent_zone_id = %s
                    ORDER BY network, id
                    """,
                    (enrollment["agent_zone_id"],),
                )
                networks = [
                    row["network"] for row in cur.fetchall()
                ]

                cur.execute(
                    """
                    INSERT INTO audit_logs (
                        actor_user_id,
                        action,
                        entity_type,
                        entity_id,
                        new_value_json,
                        created_at
                    )
                    VALUES (
                        NULL,
                        'agent.enrollment.redeemed',
                        'agent_enrollment',
                        %s,
                        %s,
                        %s
                    )
                    """,
                    (
                        enrollment["id"],
                        Jsonb(
                            {
                                "agent_id": agent_id,
                                "site_id": str(enrollment["site_id"]),
                                "agent_zone_id": str(
                                    enrollment["agent_zone_id"]
                                ),
                                "display_name": enrollment[
                                    "display_name"
                                ],
                                "used_at": used_at.isoformat(),
                            }
                        ),
                        operation_now,
                    ),
                )

                result = {
                    "agent_id": agent_id,
                    "site_id": enrollment["site_id"],
                    "site_name": zone["site_name"],
                    "agent_zone_id": enrollment["agent_zone_id"],
                    "agent_zone_name": zone["agent_zone_name"],
                    "display_name": enrollment["display_name"],
                    "ip_ranges": networks,
                }

            connection.commit()
            return result

        except UniqueViolation as exc:
            connection.rollback()
            constraint = exc.diag.constraint_name

            if constraint in {
                AGENT_ID_CONSTRAINT,
                AGENT_SECRET_HASH_CONSTRAINT,
            }:
                raise AgentEnrollmentCredentialCollision(
                    "Generated Agent credential collided"
                ) from exc

            if constraint == AGENT_ZONE_CONSTRAINT:
                raise AgentEnrollmentConflict(
                    "Agent Zone already has an assigned Agent"
                ) from exc

            raise

        except Exception:
            connection.rollback()
            raise

        finally:
            connection.close()