"""Consolidated registry database operations.

Handles Supabase DB queries for tasks, datasets, and publishing
via the Harbor registry.
"""

import ssl
from typing import Any, cast

import httpx
from pydantic import BaseModel
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from harbor.auth.client import create_authenticated_client, reset_client
from harbor.models.package.version_ref import RefType, VersionRef

RETRYABLE_EXCEPTIONS = (httpx.RequestError, ssl.SSLError, OSError)
RPC_MAX_ATTEMPTS = 3

_rpc_retry = retry(
    retry=retry_if_exception_type(RETRYABLE_EXCEPTIONS),
    stop=stop_after_attempt(RPC_MAX_ATTEMPTS),
    wait=wait_exponential(multiplier=0.5, min=0.5, max=4.0),
    before_sleep=lambda _: reset_client(),
    reraise=True,
)


def _sanitize_pg_text(value: str) -> str:
    """Strip null bytes that PostgreSQL TEXT columns cannot store."""
    return value.replace("\x00", "")


class ResolvedTaskVersion(BaseModel):
    """Result of resolving a task version reference."""

    id: str
    archive_path: str
    content_hash: str


class RegistryDB:
    # ------------------------------------------------------------------
    # Task version resolution
    # ------------------------------------------------------------------

    @_rpc_retry
    async def resolve_task_version(
        self, org: str, name: str, ref: str = "latest"
    ) -> ResolvedTaskVersion:
        """Resolve a task version reference to archive_path + content_hash.

        Handles TAG, REVISION, and DIGEST ref types via ``VersionRef.parse()``.
        """
        client = await create_authenticated_client()
        parsed = VersionRef.parse(ref)

        match parsed.type:
            case RefType.TAG:
                response = await (
                    client.table("task_version_tag")
                    .select(
                        "task_version:task_version_id(id, archive_path, content_hash), "
                        "package:package_id!inner(name, org:org_id!inner(name))"
                    )
                    .eq("tag", parsed.value)
                    .eq("package.name", name)
                    .eq("package.type", "task")
                    .eq("package.org.name", org)
                    .single()
                    .execute()
                )
                tv = cast(dict[str, Any], response.data)["task_version"]
                return ResolvedTaskVersion(
                    id=tv["id"],
                    archive_path=tv["archive_path"],
                    content_hash=tv["content_hash"],
                )

            case RefType.REVISION:
                response = await (
                    client.table("task_version")
                    .select(
                        "id, archive_path, content_hash, "
                        "package:package_id!inner(name, org:org_id!inner(name))"
                    )
                    .eq("revision", int(parsed.value))
                    .eq("package.name", name)
                    .eq("package.type", "task")
                    .eq("package.org.name", org)
                    .single()
                    .execute()
                )
                row = cast(dict[str, Any], response.data)
                return ResolvedTaskVersion(
                    id=row["id"],
                    archive_path=row["archive_path"],
                    content_hash=row["content_hash"],
                )

            case RefType.DIGEST:
                digest_value = parsed.value.removeprefix("sha256:")
                response = await (
                    client.table("task_version")
                    .select(
                        "id, archive_path, content_hash, "
                        "package:package_id!inner(name, org:org_id!inner(name))"
                    )
                    .eq("content_hash", digest_value)
                    .eq("package.name", name)
                    .eq("package.type", "task")
                    .eq("package.org.name", org)
                    .single()
                    .execute()
                )
                row = cast(dict[str, Any], response.data)
                return ResolvedTaskVersion(
                    id=row["id"],
                    archive_path=row["archive_path"],
                    content_hash=row["content_hash"],
                )

            case _:
                raise ValueError(f"Unknown ref type: {parsed.type}")

    async def resolve_task_content_hash(
        self, org: str, name: str, ref: str = "latest"
    ) -> str:
        """Convenience wrapper returning just the content_hash."""
        resolved = await self.resolve_task_version(org, name, ref)
        return resolved.content_hash

    @_rpc_retry
    async def task_version_exists(self, org: str, name: str, content_hash: str) -> bool:
        """Check whether a task version with the given content_hash already exists."""
        client = await create_authenticated_client()
        response = await (
            client.table("task_version")
            .select("id, package:package_id!inner(name, org:org_id!inner(name))")
            .eq("content_hash", content_hash)
            .eq("package.name", name)
            .eq("package.type", "task")
            .eq("package.org.name", org)
            .maybe_single()
            .execute()
        )
        return response is not None and response.data is not None

    # ------------------------------------------------------------------
    # Dataset version resolution
    # ------------------------------------------------------------------

    @_rpc_retry
    async def resolve_dataset_version(
        self, org: str, name: str, ref: str = "latest"
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Resolve a dataset version reference to (package_row, dataset_version_row)."""
        client = await create_authenticated_client()
        parsed = VersionRef.parse(ref)

        match parsed.type:
            case RefType.TAG:
                response = await (
                    client.table("dataset_version_tag")
                    .select(
                        "dataset_version:dataset_version_id(*), "
                        "package:package_id!inner(*, org:org_id!inner(name))"
                    )
                    .eq("tag", parsed.value)
                    .eq("package.name", name)
                    .eq("package.type", "dataset")
                    .eq("package.org.name", org)
                    .limit(1)
                    .execute()
                )
                data = cast(list[dict[str, Any]], response.data or [])
                if not data:
                    raise ValueError(
                        f"Tag '{parsed.value}' not found for dataset '{org}/{name}'"
                    )
                return data[0]["package"], data[0]["dataset_version"]

            case RefType.REVISION:
                response = await (
                    client.table("dataset_version")
                    .select("*, package:package_id!inner(*, org:org_id!inner(name))")
                    .eq("revision", int(parsed.value))
                    .eq("package.name", name)
                    .eq("package.type", "dataset")
                    .eq("package.org.name", org)
                    .limit(1)
                    .execute()
                )
                data = cast(list[dict[str, Any]], response.data or [])
                if not data:
                    raise ValueError(
                        f"Revision {parsed.value} not found for dataset '{org}/{name}'"
                    )
                row = data[0]
                package = row.pop("package")
                return package, row

            case RefType.DIGEST:
                response = await (
                    client.table("dataset_version")
                    .select("*, package:package_id!inner(*, org:org_id!inner(name))")
                    .eq("content_hash", parsed.value)
                    .eq("package.name", name)
                    .eq("package.type", "dataset")
                    .eq("package.org.name", org)
                    .limit(1)
                    .execute()
                )
                data = cast(list[dict[str, Any]], response.data or [])
                if not data:
                    raise ValueError(
                        f"Digest '{parsed.value}' not found for dataset '{org}/{name}'"
                    )
                row = data[0]
                package = row.pop("package")
                return package, row

            case _:
                raise ValueError(f"Unknown ref type: {parsed.type}")

    @_rpc_retry
    async def get_dataset_version_tasks(
        self, dataset_version_id: str
    ) -> list[dict[str, Any]]:
        """Return task rows for a dataset version."""
        client = await create_authenticated_client()
        response = await (
            client.table("dataset_version_task")
            .select(
                "task_version:task_version_id("
                "content_hash, "
                "package:package_id(name, org:org_id(name))"
                ")"
            )
            .eq("dataset_version_id", dataset_version_id)
            .execute()
        )
        return cast(list[dict[str, Any]], response.data or [])

    @_rpc_retry
    async def get_dataset_version_files(
        self, dataset_version_id: str
    ) -> list[dict[str, Any]]:
        """Return file rows for a dataset version."""
        client = await create_authenticated_client()
        response = await (
            client.table("dataset_version_file")
            .select("path, storage_path, content_hash")
            .eq("dataset_version_id", dataset_version_id)
            .execute()
        )
        return cast(list[dict[str, Any]], response.data or [])

    # ------------------------------------------------------------------
    # User / auth helpers
    # ------------------------------------------------------------------

    @_rpc_retry
    async def get_user_id(self) -> str:
        client = await create_authenticated_client()
        response = await client.auth.get_user()
        if response is None or response.user is None:
            raise RuntimeError(
                "Not authenticated. Please run `harbor auth login` first."
            )
        return response.user.id

    # ------------------------------------------------------------------
    # Publishing RPCs
    # ------------------------------------------------------------------

    @_rpc_retry
    async def ensure_org(self, org: str) -> dict[str, Any]:
        """Ensure an organization exists, creating it if needed."""
        client = await create_authenticated_client()
        response = await client.rpc(
            "ensure_org",
            {"p_org": org},
        ).execute()
        return cast(dict[str, Any], response.data)

    @_rpc_retry
    async def publish_task_version(
        self,
        *,
        org: str,
        name: str,
        tags: list[str],
        content_hash: str,
        archive_path: str,
        description: str | None,
        authors: list[dict[str, Any]],
        keywords: list[str],
        metadata: dict[str, Any],
        verifier_config: dict[str, Any],
        agent_config: dict[str, Any],
        environment_config: dict[str, Any],
        instruction: str,
        readme: str,
        files: list[dict[str, Any]],
        visibility: str = "public",
    ) -> dict[str, Any]:
        """Publish a task version via the publish_task_version RPC function."""
        client = await create_authenticated_client()
        response = await client.rpc(
            "publish_task_version",
            {
                "p_org": org,
                "p_name": name,
                "p_tags": tags,
                "p_content_hash": content_hash,
                "p_archive_path": archive_path,
                "p_description": _sanitize_pg_text(description)
                if description
                else description,
                "p_authors": authors,
                "p_keywords": keywords,
                "p_metadata": metadata,
                "p_verifier_config": verifier_config,
                "p_agent_config": agent_config,
                "p_environment_config": environment_config,
                "p_instruction": _sanitize_pg_text(instruction),
                "p_readme": _sanitize_pg_text(readme),
                "p_files": files,
                "p_visibility": visibility,
            },
        ).execute()
        return cast(dict[str, Any], response.data)

    @_rpc_retry
    async def publish_dataset_version(
        self,
        *,
        org: str,
        name: str,
        tags: list[str],
        description: str | None,
        authors: list[dict[str, Any]],
        tasks: list[dict[str, str]],
        files: list[dict[str, Any]],
        visibility: str = "public",
        promote_tasks: bool = False,
    ) -> dict[str, Any]:
        """Publish a dataset version via the publish_dataset_version RPC function."""
        client = await create_authenticated_client()
        response = await client.rpc(
            "publish_dataset_version",
            {
                "p_org": org,
                "p_name": name,
                "p_tags": tags,
                "p_description": _sanitize_pg_text(description)
                if description
                else description,
                "p_authors": authors,
                "p_tasks": tasks,
                "p_files": files,
                "p_visibility": visibility,
                "p_promote_tasks": promote_tasks,
            },
        ).execute()
        return cast(dict[str, Any], response.data)

    # ------------------------------------------------------------------
    # Package helpers
    # ------------------------------------------------------------------

    @_rpc_retry
    async def get_private_dataset_task_count(self, *, org: str, name: str) -> int:
        """Count private tasks linked to a dataset package."""
        client = await create_authenticated_client()
        response = await client.rpc(
            "get_private_dataset_task_count",
            {"p_org": org, "p_name": name},
        ).execute()
        return cast(int, response.data) if response.data else 0

    @_rpc_retry
    async def get_package_type(self, *, org: str, name: str) -> str | None:
        """Query the package table to get the package type (task/dataset)."""
        client = await create_authenticated_client()
        response = await (
            client.table("package")
            .select("type, org:org_id!inner(name)")
            .eq("name", name)
            .eq("org.name", org)
            .limit(1)
            .execute()
        )
        data = cast(list[dict[str, Any]], response.data or [])
        if not data:
            return None
        return data[0]["type"]

    @_rpc_retry
    async def get_package_visibility(self, *, org: str, name: str) -> str | None:
        """Query the package table to get the current visibility."""
        client = await create_authenticated_client()
        response = await (
            client.table("package")
            .select("visibility, org:org_id!inner(name)")
            .eq("name", name)
            .eq("org.name", org)
            .limit(1)
            .execute()
        )
        data = cast(list[dict[str, Any]], response.data or [])
        if not data:
            return None
        return data[0]["visibility"]

    # ------------------------------------------------------------------
    # Download recording
    # ------------------------------------------------------------------

    @_rpc_retry
    async def record_task_download(self, task_version_id: str) -> None:
        client = await create_authenticated_client()
        await (
            client.table("task_version_download")
            .insert({"task_version_id": task_version_id})
            .execute()
        )

    @_rpc_retry
    async def record_dataset_download(self, dataset_version_id: str) -> None:
        client = await create_authenticated_client()
        await (
            client.table("dataset_version_download")
            .insert({"dataset_version_id": dataset_version_id})
            .execute()
        )

    # ------------------------------------------------------------------
    # Visibility
    # ------------------------------------------------------------------

    @_rpc_retry
    async def set_package_visibility(
        self,
        *,
        org: str,
        name: str,
        package_type: str,
        visibility: str | None = None,
        toggle: bool = False,
        cascade: bool = False,
    ) -> dict[str, Any]:
        """Set or toggle visibility for a package (task or dataset)."""
        client = await create_authenticated_client()
        response = await client.rpc(
            "set_package_visibility",
            {
                "p_org": org,
                "p_name": name,
                "p_package_type": package_type,
                "p_visibility": visibility,
                "p_toggle": toggle,
                "p_cascade": cascade,
            },
        ).execute()
        return cast(dict[str, Any], response.data)
