"""PostgreSQL integration tests for secure Agent enrollment."""

import unittest
from concurrent.futures import ThreadPoolExecutor
from uuid import uuid4

from psycopg.errors import ForeignKeyViolation
from psycopg.rows import dict_row

from app.agent_auth.exceptions import AgentAuthenticationFailed
from app.agent_auth.repository import AgentAuthRepository
from app.agent_auth.service import AgentAuthService, hash_agent_secret
from app.agent_enrollment.exceptions import (
    AgentEnrollmentConflict,
    AgentEnrollmentCredentialCollision,
    AgentEnrollmentNotFound,
    AgentEnrollmentTokenInvalid,
)
from app.agent_enrollment.models import (
    AgentEnrollmentCreateRequest,
    AgentEnrollRequest,
)
from app.agent_enrollment.repository import (
    AgentEnrollmentRepository,
    default_connection_factory,
)
from app.agent_enrollment.service import AgentEnrollmentService
from app.agent_zones.models import AgentCheckinRequest
from app.agent_zones.repository import AgentZoneRepository
from app.agent_zones.service import AgentZoneService
from app.config import settings


class AgentEnrollmentPostgresTests(unittest.TestCase):
    def setUp(self):
        self.site_id = uuid4()
        self.zone_id = uuid4()
        self.site_name = f"Enrollment Test Site {self.site_id.hex[:8]}"
        self.zone_name = f"Enrollment Test Zone {self.zone_id.hex[:8]}"

        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO sites (id, name)
                    VALUES (%s, %s)
                    """,
                    (self.site_id, self.site_name),
                )
                cur.execute(
                    """
                    INSERT INTO agent_zones (
                        id,
                        site_id,
                        name
                    )
                    VALUES (%s, %s, %s)
                    """,
                    (
                        self.zone_id,
                        self.site_id,
                        self.zone_name,
                    ),
                )
                cur.execute(
                    """
                    INSERT INTO agent_zone_network_ranges (
                        agent_zone_id,
                        network
                    )
                    VALUES (%s, '10.250.0.0/24'::cidr)
                    """,
                    (self.zone_id,),
                )
            connection.commit()
        finally:
            connection.close()

        self.service = AgentEnrollmentService(
            AgentEnrollmentRepository()
        )

    def tearDown(self):
        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    DELETE FROM audit_logs
                    WHERE entity_type = 'agent_enrollment'
                      AND (
                          new_value_json ->> 'site_id' = %s
                          OR entity_id IN (
                              SELECT id
                              FROM agent_enrollments
                              WHERE site_id = %s
                          )
                      )
                    """,
                    (str(self.site_id), self.site_id),
                )
                cur.execute(
                    """
                    DELETE FROM agent_enrollments
                    WHERE site_id = %s
                    """,
                    (self.site_id,),
                )
                cur.execute(
                    """
                    DELETE FROM minerops_site_agents
                    WHERE site_id = %s
                    """,
                    (self.site_id,),
                )
                cur.execute(
                    """
                    DELETE FROM agent_zone_network_ranges
                    WHERE agent_zone_id = %s
                    """,
                    (self.zone_id,),
                )
                cur.execute(
                    """
                    DELETE FROM agent_zones
                    WHERE id = %s
                    """,
                    (self.zone_id,),
                )
                cur.execute(
                    """
                    DELETE FROM sites
                    WHERE id = %s
                    """,
                    (self.site_id,),
                )
            connection.commit()
        finally:
            connection.close()

    def create_enrollment(self, display_name="Integration Agent"):
        # created_by_user_id is nullable at the persistence layer.
        # Router-level tests separately prove that the real admin path supplies
        # the authenticated user's UUID.
        return self.service.create_enrollment(
            AgentEnrollmentCreateRequest(
                site_id=self.site_id,
                agent_zone_id=self.zone_id,
                display_name=display_name,
            ),
            actor_user_id=None,
        )

    def test_replacement_revokes_previous_open_enrollment(self):
        first = self.create_enrollment("First Agent")
        second = self.create_enrollment("Replacement Agent")

        self.assertNotEqual(
            first["enrollment_token"],
            second["enrollment_token"],
        )

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        token_hash,
                        used_at,
                        revoked_at
                    FROM agent_enrollments
                    WHERE site_id = %s
                      AND agent_zone_id = %s
                    """,
                    (self.site_id, self.zone_id),
                )
                rows = cur.fetchall()
        finally:
            connection.close()

        self.assertEqual(len(rows), 2)

        by_hash = {
            row["token_hash"]: row
            for row in rows
        }

        first_hash = hash_agent_secret(
            first["enrollment_token"]
        )
        second_hash = hash_agent_secret(
            second["enrollment_token"]
        )

        self.assertIsNotNone(by_hash[first_hash]["revoked_at"])
        self.assertIsNone(by_hash[first_hash]["used_at"])

        self.assertIsNone(by_hash[second_hash]["revoked_at"])
        self.assertIsNone(by_hash[second_hash]["used_at"])

        with self.assertRaises(AgentEnrollmentTokenInvalid):
            self.service.redeem(
                AgentEnrollRequest(
                    token=first["enrollment_token"]
                )
            )

        listed = self.service.list_enrollments(
            site_id=self.site_id,
            agent_zone_id=self.zone_id,
        )

        self.assertEqual(len(listed), 2)

        for item in listed:
            self.assertNotIn("token_hash", item)
            self.assertNotIn("agent_secret_hash", item)
            self.assertNotIn("enrollment_token", item)
            self.assertNotIn("agent_secret", item)

    def test_concurrent_creation_leaves_exactly_one_open_enrollment(self):
        def create_one(number):
            service = AgentEnrollmentService(
                AgentEnrollmentRepository()
            )
            return service.create_enrollment(
                AgentEnrollmentCreateRequest(
                    site_id=self.site_id,
                    agent_zone_id=self.zone_id,
                    display_name=f"Concurrent Agent {number}",
                ),
                actor_user_id=None,
            )

        with ThreadPoolExecutor(max_workers=2) as executor:
            results = list(
                executor.map(create_one, (1, 2))
            )

        self.assertEqual(len(results), 2)

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        count(*) AS total,
                        count(*) FILTER (
                            WHERE used_at IS NULL
                              AND revoked_at IS NULL
                        ) AS open_count,
                        count(*) FILTER (
                            WHERE revoked_at IS NOT NULL
                        ) AS revoked_count,
                        count(*) FILTER (
                            WHERE revoked_at IS NOT NULL
                              AND revoked_at < created_at
                        ) AS invalid_revocation_count
                    FROM agent_enrollments
                    WHERE site_id = %s
                      AND agent_zone_id = %s
                    """,
                    (self.site_id, self.zone_id),
                )
                counts = cur.fetchone()
        finally:
            connection.close()

        self.assertEqual(counts["total"], 2)
        self.assertEqual(counts["open_count"], 1)
        self.assertEqual(counts["revoked_count"], 1)
        self.assertEqual(counts["invalid_revocation_count"], 0)

    def test_concurrent_redemption_has_exactly_one_winner(self):
        enrollment = self.create_enrollment()
        token = enrollment["enrollment_token"]

        def redeem_once():
            service = AgentEnrollmentService(
                AgentEnrollmentRepository()
            )
            try:
                return (
                    "success",
                    service.redeem(
                        AgentEnrollRequest(token=token)
                    ),
                )
            except AgentEnrollmentTokenInvalid:
                return ("invalid", None)

        with ThreadPoolExecutor(max_workers=2) as executor:
            results = list(
                executor.map(lambda _: redeem_once(), range(2))
            )

        successes = [
            result
            for status, result in results
            if status == "success"
        ]
        failures = [
            status
            for status, _ in results
            if status == "invalid"
        ]

        self.assertEqual(len(successes), 1)
        self.assertEqual(len(failures), 1)

        credentials = successes[0]
        agent_id = credentials["agent_id"]
        agent_secret = credentials["agent_secret"]

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        agent_id,
                        agent_secret_hash,
                        credential_created_at,
                        credential_revoked_at,
                        legacy_auth_disabled_at,
                        last_payload::text AS last_payload
                    FROM minerops_site_agents
                    WHERE site_id = %s
                      AND agent_zone_id = %s
                    """,
                    (self.site_id, self.zone_id),
                )
                agents = cur.fetchall()

                cur.execute(
                    """
                    SELECT
                        used_at,
                        revoked_at,
                        enrolled_agent_id,
                        token_hash
                    FROM agent_enrollments
                    WHERE id = %s
                    """,
                    (enrollment["id"],),
                )
                enrollment_row = cur.fetchone()

                cur.execute(
                    """
                    SELECT retired_at
                    FROM agent_legacy_auth_retirements
                    WHERE agent_id = %s
                    """,
                    (agent_id,),
                )
                retirement_row = cur.fetchone()

                cur.execute(
                    """
                    SELECT new_value_json::text AS payload
                    FROM audit_logs
                    WHERE entity_type = 'agent_enrollment'
                      AND entity_id = %s
                    ORDER BY created_at
                    """,
                    (enrollment["id"],),
                )
                audit_rows = cur.fetchall()
        finally:
            connection.close()

        self.assertEqual(len(agents), 1)
        agent = agents[0]

        self.assertEqual(agent["agent_id"], agent_id)
        self.assertEqual(
            agent["agent_secret_hash"],
            hash_agent_secret(agent_secret),
        )
        self.assertIsNotNone(agent["credential_created_at"])
        self.assertIsNone(agent["credential_revoked_at"])
        self.assertEqual(
            agent["legacy_auth_disabled_at"],
            agent["credential_created_at"],
        )
        self.assertIsNotNone(retirement_row)
        self.assertEqual(
            retirement_row["retired_at"],
            agent["credential_created_at"],
        )

        self.assertIsNotNone(enrollment_row["used_at"])
        self.assertIsNone(enrollment_row["revoked_at"])
        self.assertEqual(
            enrollment_row["enrolled_agent_id"],
            agent_id,
        )

        # Neither plaintext credential may persist in operational state.
        persisted_text = "\n".join(
            [
                str(enrollment_row),
                str(agent),
                *(row["payload"] for row in audit_rows),
            ]
        )

        self.assertNotIn(token, persisted_text)
        self.assertNotIn(agent_secret, persisted_text)
        self.assertNotIn(agent_secret, agent["last_payload"])

        # The newly issued per-Agent credential must immediately authenticate.
        auth_service = AgentAuthService(
            AgentAuthRepository()
        )

        self.assertEqual(
            auth_service.authenticate(
                agent_id,
                agent_secret,
            ),
            agent_id,
        )

        with self.assertRaises(AgentAuthenticationFailed):
            auth_service.authenticate(
                agent_id,
                "definitely-not-the-issued-secret",
            )

        original_legacy_id = settings.minerops_agent_id
        original_legacy_secret = settings.minerops_agent_shared_secret
        settings.minerops_agent_id = agent_id
        settings.minerops_agent_shared_secret = "global-secret-must-not-work"
        try:
            with self.assertRaises(AgentAuthenticationFailed):
                auth_service.authenticate(None, "global-secret-must-not-work")
        finally:
            settings.minerops_agent_id = original_legacy_id
            settings.minerops_agent_shared_secret = original_legacy_secret

        # Normal check-in must never mutate the credential.
        secret_hash_before = agent["agent_secret_hash"]

        AgentZoneService(
            AgentZoneRepository()
        ).checkin(
            agent_id,
            AgentCheckinRequest(
                site_id=self.site_id,
                agent_zone_id=self.zone_id,
                site_name=self.site_name,
                hostname="integration-test-agent",
                version="integration-test",
            ),
        )

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        agent_secret_hash,
                        credential_created_at,
                        credential_revoked_at,
                        legacy_auth_disabled_at
                    FROM minerops_site_agents
                    WHERE agent_id = %s
                    """,
                    (agent_id,),
                )
                after_checkin = cur.fetchone()
        finally:
            connection.close()

        self.assertEqual(
            after_checkin["agent_secret_hash"],
            secret_hash_before,
        )
        self.assertIsNotNone(
            after_checkin["credential_created_at"]
        )
        self.assertIsNone(
            after_checkin["credential_revoked_at"]
        )
        self.assertEqual(
            after_checkin["legacy_auth_disabled_at"],
            agent["legacy_auth_disabled_at"],
        )

        # 0017 enrollment integrity remains RESTRICT: the current Agent row
        # cannot be deleted while its enrollment relationship exists.
        connection = default_connection_factory()
        try:
            with self.assertRaises(ForeignKeyViolation):
                with connection.cursor() as cur:
                    cur.execute(
                        "DELETE FROM minerops_site_agents WHERE agent_id = %s",
                        (agent_id,),
                    )
            connection.rollback()
        finally:
            connection.close()

        # After an operator deliberately removes the enrollment history, the
        # independent retirement ledger must outlive deletion and recreation
        # of the current Agent row with the identical identity.
        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    "DELETE FROM agent_enrollments WHERE id = %s",
                    (enrollment["id"],),
                )
                cur.execute(
                    "DELETE FROM minerops_site_agents WHERE agent_id = %s",
                    (agent_id,),
                )
            connection.commit()
        finally:
            connection.close()

        original_legacy_id = settings.minerops_agent_id
        original_legacy_secret = settings.minerops_agent_shared_secret
        settings.minerops_agent_id = agent_id
        settings.minerops_agent_shared_secret = "recreation-global-secret"
        try:
            with self.assertRaises(AgentAuthenticationFailed):
                auth_service.authenticate(None, "recreation-global-secret")

            connection = default_connection_factory()
            try:
                with connection.cursor(row_factory=dict_row) as cur:
                    cur.execute(
                        """
                        INSERT INTO minerops_site_agents (
                            agent_id,
                            site_id,
                            agent_zone_id,
                            site_name
                        )
                        VALUES (%s, %s, %s, %s)
                        RETURNING legacy_auth_disabled_at
                        """,
                        (
                            agent_id,
                            self.site_id,
                            self.zone_id,
                            self.site_name,
                        ),
                    )
                    recreated = cur.fetchone()
                connection.commit()
            finally:
                connection.close()

            self.assertEqual(
                recreated["legacy_auth_disabled_at"],
                retirement_row["retired_at"],
            )
            with self.assertRaises(AgentAuthenticationFailed):
                auth_service.authenticate(agent_id, "recreation-global-secret")
        finally:
            settings.minerops_agent_id = original_legacy_id
            settings.minerops_agent_shared_secret = original_legacy_secret

    def test_failed_credential_issuance_rolls_back_new_retirement(self):
        enrollment = self.create_enrollment("Collision Agent")
        collision_id = f"agent-collision-{uuid4().hex}"

        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO minerops_site_agents (agent_id, site_name)
                    VALUES (%s, 'Existing uncredentialed identity')
                    """,
                    (collision_id,),
                )
            connection.commit()
        finally:
            connection.close()

        try:
            with self.assertRaises(AgentEnrollmentCredentialCollision):
                AgentEnrollmentRepository().redeem_enrollment(
                    token_hash=hash_agent_secret(enrollment["enrollment_token"]),
                    agent_id=collision_id,
                    agent_secret_hash=hash_agent_secret("never-issued-secret"),
                )

            connection = default_connection_factory()
            try:
                with connection.cursor(row_factory=dict_row) as cur:
                    cur.execute(
                        """
                        SELECT count(*) AS count
                        FROM agent_legacy_auth_retirements
                        WHERE agent_id = %s
                        """,
                        (collision_id,),
                    )
                    self.assertEqual(cur.fetchone()["count"], 0)
                    cur.execute(
                        """
                        SELECT used_at, enrolled_agent_id
                        FROM agent_enrollments
                        WHERE id = %s
                        """,
                        (enrollment["id"],),
                    )
                    enrollment_row = cur.fetchone()
            finally:
                connection.close()

            self.assertIsNone(enrollment_row["used_at"])
            self.assertIsNone(enrollment_row["enrolled_agent_id"])
        finally:
            connection = default_connection_factory()
            try:
                with connection.cursor() as cur:
                    cur.execute(
                        "DELETE FROM minerops_site_agents WHERE agent_id = %s",
                        (collision_id,),
                    )
                connection.commit()
            finally:
                connection.close()

    def test_expired_token_is_rejected_without_side_effects(self):
        enrollment = self.create_enrollment()

        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    UPDATE agent_enrollments
                    SET created_at = now() - interval '2 hours',
                        expires_at = now() - interval '1 hour'
                    WHERE id = %s
                    """,
                    (enrollment["id"],),
                )
            connection.commit()
        finally:
            connection.close()

        with self.assertRaises(AgentEnrollmentTokenInvalid):
            self.service.redeem(
                AgentEnrollRequest(
                    token=enrollment["enrollment_token"]
                )
            )

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT used_at, revoked_at, enrolled_agent_id
                    FROM agent_enrollments
                    WHERE id = %s
                    """,
                    (enrollment["id"],),
                )
                row = cur.fetchone()

                cur.execute(
                    """
                    SELECT count(*) AS count
                    FROM minerops_site_agents
                    WHERE agent_zone_id = %s
                    """,
                    (self.zone_id,),
                )
                agent_count = cur.fetchone()["count"]
        finally:
            connection.close()

        self.assertIsNone(row["used_at"])
        self.assertIsNone(row["revoked_at"])
        self.assertIsNone(row["enrolled_agent_id"])
        self.assertEqual(agent_count, 0)

    def test_disabled_and_occupied_zone_do_not_consume_token(self):
        enrollment = self.create_enrollment()
        token = enrollment["enrollment_token"]

        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    UPDATE agent_zones
                    SET is_enabled = false
                    WHERE id = %s
                    """,
                    (self.zone_id,),
                )
            connection.commit()
        finally:
            connection.close()

        with self.assertRaises(AgentEnrollmentConflict):
            self.service.redeem(
                AgentEnrollRequest(token=token)
            )

        connection = default_connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    UPDATE agent_zones
                    SET is_enabled = true
                    WHERE id = %s
                    """,
                    (self.zone_id,),
                )
                cur.execute(
                    """
                    INSERT INTO minerops_site_agents (
                        agent_id,
                        site_id,
                        agent_zone_id,
                        site_name
                    )
                    VALUES (%s, %s, %s, %s)
                    """,
                    (
                        f"legacy-{uuid4().hex}",
                        self.site_id,
                        self.zone_id,
                        self.site_name,
                    ),
                )
            connection.commit()
        finally:
            connection.close()

        with self.assertRaises(AgentEnrollmentConflict):
            self.service.redeem(
                AgentEnrollRequest(token=token)
            )

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT used_at, revoked_at, enrolled_agent_id
                    FROM agent_enrollments
                    WHERE id = %s
                    """,
                    (enrollment["id"],),
                )
                row = cur.fetchone()
        finally:
            connection.close()

        self.assertIsNone(row["used_at"])
        self.assertIsNone(row["revoked_at"])
        self.assertIsNone(row["enrolled_agent_id"])

    def test_site_zone_mismatch_is_rejected_before_creation(self):
        wrong_site_id = uuid4()

        with self.assertRaises(AgentEnrollmentNotFound):
            self.service.create_enrollment(
                AgentEnrollmentCreateRequest(
                    site_id=wrong_site_id,
                    agent_zone_id=self.zone_id,
                    display_name="Wrong Site Agent",
                ),
                actor_user_id=None,
            )

        connection = default_connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT count(*) AS count
                    FROM agent_enrollments
                    WHERE agent_zone_id = %s
                    """,
                    (self.zone_id,),
                )
                count = cur.fetchone()["count"]
        finally:
            connection.close()

        self.assertEqual(count, 0)


if __name__ == "__main__":
    unittest.main()