from __future__ import annotations

import json
import subprocess
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path

from ingestion.common.classification import classify_table
from ingestion.common.gcp import list_objects
from ingestion.common.models import table_to_manifest

from .config import SourceConfig
from .source_db import inspect_tables


def write_manifest(config: SourceConfig, run_date: str, run_id: str, upload: bool = True) -> Path:
    gcs_uri = config.gcs_run_uri(run_date, run_id)
    objects = list_gcs_objects(gcs_uri)
    manifest_uri = config.manifest_gcs_uri(run_date, run_id)
    manifest = build_manifest(config, run_date, run_id, gcs_uri, manifest_uri, objects)

    output_path = config.root / "runs" / run_date / run_id / "manifest.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    if upload:
        subprocess.run(
            ["gcloud", "storage", "cp", str(output_path), manifest_uri],
            check=True,
            capture_output=True,
            text=True,
        )
    return output_path


def list_gcs_objects(gcs_uri: str) -> list[dict]:
    """List objects under a GCS prefix as ``{key, size}`` (``key`` is the object path
    relative to the bucket, matching the S3 manifest's ``files[].key`` shape)."""
    return [{"key": item["name"], "size": item["size"]} for item in list_objects(gcs_uri)]


def build_manifest(
    config: SourceConfig,
    run_date: str,
    run_id: str,
    gcs_uri: str,
    manifest_uri: str,
    objects: list[dict],
) -> dict:
    table_metadata = load_table_metadata(config)
    summaries = summarize_tables(config, run_date, run_id, objects, table_metadata)
    return {
        "manifest_schema_version": 1,
        "source_name": config.source_name,
        "project": config.project,
        "source_database": config.dataset,
        "dataset": config.dataset,
        "schema_source": "information_schema",
        "run_date": run_date,
        "run_id": run_id,
        "product_status": "exported",
        "gcs_uri": gcs_uri,
        "manifest_uri": manifest_uri,
        "runtime_log_uri": config.runtime_log_gcs_uri(run_date, run_id),
        "created_at": datetime.now(UTC).isoformat(),
        "objects": {
            "file_count": len(objects),
            "total_bytes": sum(item["size"] for item in objects),
            "files": objects,
        },
        "tables": summaries,
        "schema_drift": schema_drift(
            exported=(str(item["table_name"]) for item in summaries),
            expected=table_metadata.keys(),
        ),
    }


def load_table_metadata(config: SourceConfig) -> dict[str, dict]:
    rules = config.classification_rules
    return {
        table.name: table_to_manifest(table, classify_table(table, rules))
        for table in inspect_tables(config)
    }


def summarize_tables(
    config: SourceConfig,
    run_date: str,
    run_id: str,
    objects: list[dict],
    table_metadata: dict[str, dict],
) -> list[dict]:
    grouped: dict[str, dict] = defaultdict(lambda: {"file_count": 0, "total_bytes": 0})
    for item in objects:
        name = item["key"]
        if not name.endswith(".parquet"):
            continue
        table = table_from_key(name)
        if table is None:
            continue
        grouped[table]["table_name"] = table
        grouped[table]["file_count"] += 1
        grouped[table]["total_bytes"] += int(item["size"])
    summaries = []
    for table, summary in grouped.items():
        summaries.append(
            {
                **table_metadata.get(table, {}),
                **summary,
                "gcs_uri": config.table_gcs_uri(table, run_date, run_id),
            }
        )
    return sorted(summaries, key=lambda item: str(item["table_name"]))


def table_from_key(key: str) -> str | None:
    """Recover the table name from a ``.../table=<name>/<file>`` object key."""
    for segment in key.split("/"):
        if segment.startswith("table="):
            return segment[len("table=") :]
    return None


def schema_drift(exported, expected) -> dict[str, list[str]]:
    """Differences between the tables actually exported (parsed from GCS objects) and the
    tables we have schema metadata/contracts for, surfaced so a desync between export and
    contracts is observable at produce time rather than deferred to a loader hard-fail."""
    exported_set = set(exported)
    expected_set = set(expected)
    return {
        "files_without_metadata": sorted(exported_set - expected_set),
        "metadata_without_files": sorted(expected_set - exported_set),
    }