"""Classify procedures from CUR dependency closures."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

LOGGER = logging.getLogger(__name__)

WRITE_RELATIONS: frozenset[str] = frozenset(
    r.upper() for r in ("INSERT", "UPDATE", "DELETE", "MERGE", "TRUNCATE", "ALTER", "DROP")
)


@dataclass(frozen=True)
class ClassificationResult:
    """Classification output for one procedure."""
    modifies_data: bool = False
    affected_ids: frozenset[str] = field(default_factory=frozenset)
    read_ids: frozenset[str] = field(default_factory=frozenset)


def classify(
    code_unit_id: str,
    index: dict[str, object],
) -> ClassificationResult:
    """Scan a CUR dependency closure for write and read operations.

    Iterates every dependency edge in *index* (the transitive closure
    returned by ``FindOptions(include_dependencies=True)``).  If any edge
    carries a write relation type (INSERT, UPDATE, DELETE, MERGE, TRUNCATE,
    ALTER, DROP) the procedure is classified as read-write and the target
    code unit ID is collected.  Edges that carry only non-write relation
    types are collected as read dependencies.

    Returns ``modifies_data`` flag, the set of affected (write) IDs, and
    the set of read-only dependency IDs.
    """
    if code_unit_id not in index:
        return ClassificationResult()

    all_deps = [dependency for unit in index.values() for dependency in _get_deps(unit)]
    write_deps = [dependency for dependency in all_deps if _has_write_relation(dependency)]
    read_deps = [dependency for dependency in all_deps if not _has_write_relation(dependency) and _has_any_relation(dependency)]
    affected = frozenset(getattr(dependency, "id") for dependency in write_deps if getattr(dependency, "id", None))
    read_ids = frozenset(getattr(dependency, "id") for dependency in read_deps if getattr(dependency, "id", None))

    return ClassificationResult(
        modifies_data=bool(write_deps),
        affected_ids=affected,
        read_ids=read_ids,
    )


def _get_deps(unit: object) -> list:
    deps = getattr(unit, "dependencies", None)
    if deps is None:
        return []
    return getattr(deps, "dependsOn", None) or []


def _has_write_relation(dep: object) -> bool:
    types = getattr(dep, "relationTypes", None) or []
    return bool(WRITE_RELATIONS & {t.upper() for t in types})


def _has_any_relation(dep: object) -> bool:
    types = getattr(dep, "relationTypes", None) or []
    return bool(types)


def _format_name(unit: object) -> Optional[str]:
    source = getattr(unit, "source", None)
    if source is None:
        return None
    name = getattr(source, "name", None)
    if not name:
        return None
    schema = getattr(source, "schema_", None) or getattr(source, "schema", None)
    database = getattr(source, "database", None)
    parts = [p for p in (database, schema, name) if p]
    return ".".join(parts)


def _open_registry(project_root: str | Path) -> object | None:
    try:
        from snowflake_code_unit_registry import CodeUnitRegistry
    except ImportError:
        LOGGER.warning("snowflake-code-unit-registry package not installed; skipping CUR enrichment")
        return None
    try:
        return CodeUnitRegistry.open(str(project_root))
    except Exception as exc:
        LOGGER.warning("Registry not found at %s: %s", project_root, exc)
        return None


def _build_closure_index(registry: object, code_unit_id: str) -> dict[str, object]:
    from snowflake_code_unit_registry import FindOptions
    escaped_id = code_unit_id.replace("'", "''")
    units = registry.find_all(FindOptions(
        filter=f"id = '{escaped_id}'",
        include_dependencies=True,
        fields=["id", "source", "target", "dependencies", "files"],
    ))
    return {u.id: u for u in units if u.id}


def _compute_checksums(
    unit: object,
    project_root: Path,
) -> tuple[Optional[str], Optional[str]]:
    """Compute BLAKE3 checksums for the source and converted files on disk.

    Paths stored in CUR are relative to *project_root*.  Returns
    ``(source_checksum, converted_checksum)``; either may be ``None`` if the
    file path is absent or the file cannot be read.
    """
    from snowflake_code_unit_registry import compute_file_checksum

    def _checksum(file_entry: object) -> Optional[str]:
        path_str = getattr(file_entry, "path", None)
        if not path_str:
            return None
        abs_path = project_root / path_str
        try:
            return compute_file_checksum(str(abs_path))
        except Exception as exc:
            LOGGER.warning("Could not compute checksum for %s: %s", abs_path, exc)
            return None

    files = getattr(unit, "files", None)
    if files is None:
        return None, None
    return (
        _checksum(getattr(files, "source", None)),
        _checksum(getattr(files, "converted", None)),
    )


def extract_code_unit_id(test_file_path: str) -> Optional[str]:
    """Extract code unit ID from a generated test YAML path."""
    return Path(test_file_path).stem


def enrich_test_files(
    test_files: list,
    project_root: str | Path,
) -> None:
    """Populate modifies_data and affected_tables from CUR source metadata.

    Skips test files where modifies_data is already set (YAML override).
    """
    registry = _open_registry(project_root)
    if registry is None:
        return

    project_root = Path(project_root)
    for test_file in test_files:
        cu_id = extract_code_unit_id(test_file.file_path)
        if cu_id is None:
            continue
        try:
            index = _build_closure_index(registry, cu_id)
        except Exception as exc:
            LOGGER.warning("Could not query code unit %s from CUR: %s", cu_id, exc)
            continue
        if cu_id not in index:
            continue

        test_file.source_checksum, test_file.converted_checksum = _compute_checksums(
            index[cu_id], project_root,
        )

        target_obj = getattr(index[cu_id], "target", None)
        if target_obj:
            test_file.target_database = getattr(target_obj, "database", None)

        if test_file.modifies_data is not None:
            continue
        result = classify(cu_id, index)
        test_file.modifies_data = result.modifies_data
        test_file.affected_tables = sorted(
            filter(None, (_format_name(index[aid]) for aid in result.affected_ids if aid in index))
        )
        test_file.read_tables = sorted(
            filter(None, (_format_name(index[rid]) for rid in result.read_ids if rid in index))
        )
