from __future__ import annotations

from pathlib import Path
from typing import Any

import yaml

from ingestion.common.classification import classify_table
from ingestion.common.models import Table

from .config import SourceConfig
from .source_db import inspect_tables

AUTHORED_SCHEMA_VERSION = "0.1.0"
PRODUCT_VERSION = "0.1.0"


def generate_contracts(config: SourceConfig) -> list[Path]:
    """Generate the S3 snapshot *source* contracts this extract stage produces.

    The extract stage writes Parquet snapshots to S3 and owns their source contract.
    The raw-snapshot *product* contract is owned and generated by the load stage that
    materializes the Cloud Lake raw tables; staging and above are owned by the transform
    layer (see ADR 0015 and ADR 0016).
    """
    tables = inspect_tables(config)
    output_dir = config.root / config.data.get("contracts", {}).get("output_dir", "contracts/generated")
    output_dir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []
    for table in tables:
        paths.append(
            write_yaml(
                output_dir / f"{config.source_name}.{table.name}.s3-snapshot.source.contract.yaml",
                source_contract(config, table),
            )
        )
    return paths


def source_contract(config: SourceConfig, table: Table) -> dict[str, Any]:
    table_slug = table.name.lower()
    classifications = classify_table(table, config.classification_rules)
    return {
        "schema_version": AUTHORED_SCHEMA_VERSION,
        "kind": "source",
        "id": f"source.{config.source_name}.{table_slug}_s3_snapshot",
        "name": f"{config.source_name} {table.name} S3 snapshot export",
        "version": PRODUCT_VERSION,
        "owner": _owner(config),
        "domain": _domain(config),
        "purpose": (
            f"Source-faithful per-snapshot Parquet export of {config.source_database}.{table.name} "
            "from TiDB Cloud Essential to S3, consumed by the lakehouse load stage."
        ),
        "usage": {
            "intended": [
                "Bulk load into the raw lakehouse layer",
                "Audit and provenance of exported snapshots",
            ],
            "prohibited": [
                "Direct analytical querying (load into a contracted product first).",
                "Treating a single snapshot as an intra-day event stream.",
            ],
        },
        "source": _source(config, table),
        "assets": [
            {
                "id": f"{table_slug}_s3_snapshot",
                "kind": "file",
                "ref": f"s3://{config.data_bucket}/{config.snapshot_prefix}/",
                "materialization": "external",
                "produced_by": {"kind": "python", "ref": "ingestion.extract.tidbcloud.engine.export"},
            }
        ],
        "schema": _schema(table),
        "freshness": _freshness(config),
        "tests": _placeholder_tests(table_slug, "source"),
        "governance": _governance_source(config, classifications),
        "agent_permissions": {"default": "deny"},
        "compatibility": {"change_policy": "semver"},
        "extensions": {
            "primary_key": table.primary_key,
            "watermark_columns": table.watermark_columns,
            "row_estimate": table.rows_estimate,
            "known_blind_spots": [
                "Intra-snapshot changes are collapsed.",
                "Insert-then-delete between snapshots is invisible.",
                "Schema drift requires review before promotion.",
            ],
        },
    }


def _owner(config: SourceConfig) -> dict[str, str]:
    return {"type": "team", "name": config.data["source"]["owner"]}


def _domain(config: SourceConfig) -> dict[str, str]:
    return {"name": config.data["source"]["domain"]}


def _source(config: SourceConfig, table: Table) -> dict[str, Any]:
    tidb = config.data["source"].get("tidb_cloud", {})
    identifiers: dict[str, Any] = {
        "database": config.source_database,
        "table": table.name,
        "cluster_id": config.tidb_cluster_id,
    }
    if tidb.get("branch_id"):
        identifiers["branch_id"] = tidb["branch_id"]
    if tidb.get("branch_name"):
        identifiers["branch_name"] = tidb["branch_name"]
    return {
        "kind": "database",
        "system": "tidbcloud",
        "description": f"TiDB Cloud Essential {config.source_database}.{table.name}",
        "identifiers": identifiers,
    }


def _schema(table: Table) -> dict[str, Any]:
    return {
        "kind": "relational",
        "columns": [
            {
                "name": column.name.lower(),
                "type": column.column_type,
                "nullable": column.nullable,
            }
            for column in table.columns
        ],
    }


def _freshness(config: SourceConfig) -> dict[str, str]:
    freshness = config.data.get("freshness", {})
    blocked = freshness.get("blocked_after_hours", 48)
    degraded = freshness.get("degraded_after_hours", 30)
    return {"max_age": f"PT{blocked}H", "degrade_after": f"PT{degraded}H"}


def _governance_source(config: SourceConfig, classifications: list[str]) -> dict[str, Any]:
    retention_days = config.data.get("retention", {}).get("s3_snapshots_days", 90)
    return {
        "classifications": classifications,
        "access": {"default": "restricted"},
        "retention": {"period": f"P{retention_days}D"},
        "audit": {"required": True, "events": ["export", "agent_action"]},
    }


def _placeholder_tests(table_slug: str, layer: str) -> list[dict[str, Any]]:
    return [
        {
            "id": f"{table_slug}_{layer}_smoke",
            "type": "sql",
            "severity": "informational",
            "description": (
                f"Placeholder informational smoke check for the {layer} contract of {table_slug}. "
                "Replace with concrete row-count, freshness, or key-uniqueness assertions."
            ),
            "sql": "select 1",
        }
    ]


def write_yaml(path: Path, data: dict[str, Any]) -> Path:
    path.write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=False))
    return path