from __future__ import annotations

import tempfile
import unittest
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from arr_ingestion.contracts import IngestionError
from arr_storage.filesystem import FilesystemObjectBackend
from arr_storage.postgres_grants import (
    PostgresReadGrantBroker,
    ReadGrantDatabaseConfig,
)
from arr_storage.store import ManagedObjectStore
from arr_web.run import _parser


NOW = datetime(2026, 7, 28, 12, 0, tzinfo=timezone.utc)


class FakeGrantDatabase:
    def __init__(self, *, table_exists: bool = True, database: str = "booking_test") -> None:
        self.table_exists = table_exists
        self.database = database
        self.grants: dict[str, dict[str, Any]] = {}
        self.connections: list[FakeGrantConnection] = []

    def connect(self, _dsn: str) -> "FakeGrantConnection":
        connection = FakeGrantConnection(self)
        self.connections.append(connection)
        return connection


class FakeGrantConnection:
    def __init__(self, database: FakeGrantDatabase) -> None:
        self.database = database
        self.next_row: tuple[Any, ...] | None = None
        self.rowcount = -1
        self.closed = False

    def transaction(self) -> "FakeGrantConnection":
        return self

    def cursor(self) -> "FakeGrantConnection":
        return self

    def __enter__(self) -> "FakeGrantConnection":
        return self

    def __exit__(self, _kind: object, _value: object, _traceback: object) -> bool:
        return False

    def execute(self, query: str, params: tuple[Any, ...] | None = None) -> None:
        normalized = " ".join(query.split())
        self.next_row = None
        self.rowcount = -1
        if "SELECT current_database()" in normalized:
            self.next_row = (self.database.database,)
        elif "to_regclass('ingestion.source_read_grants')" in normalized:
            self.next_row = (
                "ingestion.source_read_grants" if self.database.table_exists else None,
            )
        elif normalized.startswith("INSERT INTO ingestion.source_read_grants"):
            assert params is not None
            handle, job_id, attempt_no, object_key, filename, created_at, expires_at = params
            self.database.grants[str(handle)] = {
                "job_id": str(job_id),
                "attempt_no": int(attempt_no),
                "object_key": str(object_key),
                "filename": str(filename),
                "created_at": created_at,
                "expires_at": expires_at,
                "consumed_at": None,
            }
            self.rowcount = 1
        elif normalized.startswith("UPDATE ingestion.source_read_grants"):
            assert params is not None
            consumed_at, handle, job_id, current_time = params
            grant = self.database.grants.get(str(handle))
            if (
                grant is not None
                and grant["job_id"] == job_id
                and grant["consumed_at"] is None
                and grant["expires_at"] > current_time
            ):
                grant["consumed_at"] = consumed_at
                self.next_row = (grant["object_key"], grant["filename"])
                self.rowcount = 1
            else:
                self.rowcount = 0
        elif normalized.startswith("DELETE FROM ingestion.source_read_grants"):
            assert params is not None
            current_time = params[0]
            expired = [
                handle
                for handle, grant in self.database.grants.items()
                if grant["consumed_at"] is not None
                or grant["expires_at"] <= current_time
            ]
            for handle in expired:
                del self.database.grants[handle]
            self.rowcount = len(expired)

    def fetchone(self) -> tuple[Any, ...] | None:
        return self.next_row

    def close(self) -> None:
        self.closed = True


class PostgresReadGrantBrokerTests(unittest.TestCase):
    def _store(self, root: Path) -> tuple[ManagedObjectStore, str]:
        store = ManagedObjectStore(
            FilesystemObjectBackend(root / "objects", create=True)
        )
        source = root / "source.xml"
        source.write_bytes(b"<root><row /></root>")
        stored = store.upload_committed(
            job_id="arrjob-grant-postgres-001",
            attempt_no=1,
            role="source_xml",
            source=source,
            original_filename="source.xml",
        )
        return store, stored.object_key

    def test_issue_and_atomic_single_redemption_are_job_bound(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            store, object_key = self._store(root)
            database = FakeGrantDatabase()
            broker = PostgresReadGrantBroker(
                store,
                ReadGrantDatabaseConfig("postgresql://synthetic"),
                connect=database.connect,
                clock=lambda: NOW,
            )
            broker.assert_ready()
            grant = broker.issue(
                job_id="arrjob-grant-postgres-001",
                object_key=object_key,
                original_filename="source.xml",
                ttl_seconds=300,
            )
            self.assertNotIn(object_key, repr(grant))
            self.assertEqual(grant.max_reads, 1)
            self.assertEqual(
                database.grants[grant.handle]["attempt_no"],
                1,
            )

            with self.assertRaises(IngestionError):
                broker.materialize(
                    handle=grant.handle,
                    job_id="arrjob-other",
                    destination=root / "wrong.xml",
                    max_bytes=1024,
                )
            destination = root / "runtime" / "source.xml"
            broker.materialize(
                handle=grant.handle,
                job_id="arrjob-grant-postgres-001",
                destination=destination,
                max_bytes=1024,
            )
            self.assertEqual(destination.read_bytes(), b"<root><row /></root>")
            with self.assertRaises(IngestionError):
                broker.materialize(
                    handle=grant.handle,
                    job_id="arrjob-grant-postgres-001",
                    destination=root / "second.xml",
                    max_bytes=1024,
                )
            self.assertEqual(broker.purge_expired(), 1)
            self.assertFalse(database.grants)
            self.assertTrue(all(item.closed for item in database.connections))

    def test_expiry_ttl_and_missing_schema_fail_closed(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            store, object_key = self._store(root)
            current = [NOW]
            database = FakeGrantDatabase()
            broker = PostgresReadGrantBroker(
                store,
                ReadGrantDatabaseConfig("postgresql://synthetic"),
                connect=database.connect,
                clock=lambda: current[0],
            )
            with self.assertRaises(IngestionError):
                broker.issue(
                    job_id="arrjob-grant-postgres-001",
                    object_key=object_key,
                    original_filename="source.xml",
                    ttl_seconds=301,
                )
            grant = broker.issue(
                job_id="arrjob-grant-postgres-001",
                object_key=object_key,
                original_filename="source.xml",
                ttl_seconds=60,
            )
            current[0] = NOW + timedelta(seconds=60)
            with self.assertRaises(IngestionError):
                broker.materialize(
                    handle=grant.handle,
                    job_id="arrjob-grant-postgres-001",
                    destination=root / "expired.xml",
                    max_bytes=1024,
                )
            self.assertEqual(broker.purge_expired(), 1)

            unavailable = PostgresReadGrantBroker(
                store,
                ReadGrantDatabaseConfig("postgresql://synthetic"),
                connect=FakeGrantDatabase(table_exists=False).connect,
                clock=lambda: NOW,
            )
            with self.assertRaisesRegex(IngestionError, "unavailable"):
                unavailable.assert_ready()

    def test_target_guard_migration_and_processing_flag(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            store, _object_key = self._store(root)
            broker = PostgresReadGrantBroker(
                store,
                ReadGrantDatabaseConfig("postgresql://synthetic"),
                connect=FakeGrantDatabase(database="postgres").connect,
                clock=lambda: NOW,
            )
            with self.assertRaisesRegex(IngestionError, "target"):
                broker.assert_ready()

        migration = (
            Path(__file__).resolve().parents[1]
            / "database"
            / "009_source_read_grants.sql"
        ).read_text(encoding="utf-8")
        self.assertIn("current_database() <> 'booking_test'", migration)
        self.assertIn("CREATE TABLE ingestion.source_read_grants", migration)
        self.assertIn("interval '5 minutes'", migration)
        self.assertIn("consumed_at", migration)
        self.assertNotIn("DROP SCHEMA", migration)
        self.assertFalse(_parser().parse_args([]).enable_processing)
        self.assertTrue(
            _parser().parse_args(["--enable-processing"]).enable_processing
        )


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