from __future__ import annotations

import json
import subprocess
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path
from urllib.parse import urlparse

from pymysql.err import MySQLError

from ingestion.common.aws import aws_env
from ingestion.common.classification import classify_table
from ingestion.common.models import Column, Table, table_to_manifest

from .config import SourceConfig
from .export import ensure_profile, find_ticloud_bin
from .source_db import inspect_tables


def write_manifest(config: SourceConfig, run_date: str, run_id: str, export_id: str, upload: bool = True) -> Path:
    ensure_profile(config)
    export = find_export(config, export_id)
    s3_uri = export["target"]["s3"]["uri"]
    objects = list_s3_objects(config, s3_uri)
    manifest_uri = config.manifest_s3_uri(run_date, run_id)
    manifest = build_manifest(config, run_date, run_id, export, s3_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(
            ["aws", "s3", "cp", str(output_path), manifest_uri, "--only-show-errors"],
            check=True,
            env=aws_env(config),
        )
    return output_path


def find_export(config: SourceConfig, export_id: str) -> dict:
    payload = run_json(
        [
            find_ticloud_bin(),
            "serverless",
            "export",
            "list",
            "--profile",
            config.tidb_cli_profile,
            "--cluster-id",
            config.tidb_cluster_id,
            "--output",
            "json",
        ]
    )
    for export in payload.get("exports", []):
        if export.get("exportId") == export_id:
            return export
    raise SystemExit(f"export {export_id} was not found for cluster {config.tidb_cluster_id}")


def list_s3_objects(config: SourceConfig, s3_uri: str) -> list[dict]:
    parsed = parse_s3_uri(s3_uri)
    payload = run_json(
        [
            "aws",
            "s3api",
            "list-objects-v2",
            "--bucket",
            parsed["bucket"],
            "--prefix",
            parsed["prefix"],
            "--output",
            "json",
        ],
        env=aws_env(config),
    )
    return payload.get("Contents", [])


def build_manifest(
    config: SourceConfig,
    run_date: str,
    run_id: str,
    export: dict,
    s3_uri: str,
    manifest_uri: str,
    objects: list[dict],
) -> dict:
    files = [
        {"key": item["Key"], "size": item["Size"], "last_modified": item["LastModified"]}
        for item in objects
    ]
    table_metadata, schema_source = load_table_metadata(config)
    summaries = summarize_tables(files, table_metadata)
    return {
        "manifest_schema_version": 1,
        "source_name": config.source_name,
        "cluster_id": config.tidb_cluster_id,
        "source_database": config.source_database,
        "schema_source": schema_source,
        "source_branch_id": config.data["source"]["tidb_cloud"].get("branch_id", ""),
        "source_branch_name": config.data["source"]["tidb_cloud"].get("branch_name", ""),
        "run_date": run_date,
        "run_id": run_id,
        "export_task_id": export["exportId"],
        "export_state": export["state"],
        "product_status": product_status_for(export["state"]),
        "s3_uri": s3_uri,
        "manifest_uri": manifest_uri,
        "runtime_log_uri": config.runtime_log_s3_uri(run_date, run_id),
        "created_at": datetime.now(UTC).isoformat(),
        "export": sanitized_export(export),
        "objects": {
            "file_count": len(files),
            "total_bytes": sum(item["size"] for item in files),
            "files": files,
        },
        "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) -> tuple[dict[str, dict], str]:
    try:
        tables = inspect_tables(config)
        rules = config.classification_rules
        return {
            table.name: table_to_manifest(table, classify_table(table, rules))
            for table in tables
        }, "information_schema"
    except MySQLError:
        tables = load_tables_from_contracts(config)
        if not tables:
            raise
        rules = config.classification_rules
        return {
            table.name: table_to_manifest(table, table.classifications or classify_table(table, rules))
            for table in tables
        }, "generated_contracts"


def load_tables_from_contracts(config: SourceConfig) -> list[Table]:
    """Best-effort schema fallback from the per-table S3 snapshot source contracts.

    Used only when live information_schema is unavailable. NOTE: the authored contract
    schema requires lowercase column identifiers, so column names recovered here are
    contract-cased, not source-faithful. Live information_schema is authoritative for
    case-sensitive COPY; treat this path as degraded.
    """
    import yaml

    output_dir = config.root / config.data.get("contracts", {}).get("output_dir", "contracts/generated")
    tables = []
    for contract_path in sorted(output_dir.glob(f"{config.source_name}.*.s3-snapshot.source.contract.yaml")):
        contract = yaml.safe_load(contract_path.read_text())
        identifiers = contract.get("source", {}).get("identifiers", {})
        columns = contract.get("schema", {}).get("columns", [])
        extensions = contract.get("extensions", {})
        tables.append(
            Table(
                database=str(identifiers.get("database", config.source_database)),
                name=str(identifiers.get("table")),
                rows_estimate=extensions.get("row_estimate"),
                columns=[
                    Column(
                        name=column["name"],
                        data_type=str(column["type"]).split("(")[0].lower(),
                        column_type=column["type"],
                        nullable=bool(column.get("nullable", True)),
                        ordinal_position=index,
                    )
                    for index, column in enumerate(columns, start=1)
                ],
                primary_key=list(extensions.get("primary_key", [])),
                classifications=list(contract.get("governance", {}).get("classifications", [])),
                watermark_candidates=list(extensions.get("watermark_columns", [])),
            )
        )
    return tables


def sanitized_export(export: dict) -> dict:
    allowed = {
        "clusterId",
        "completeTime",
        "createTime",
        "displayName",
        "exportId",
        "exportOptions",
        "name",
        "snapshotTime",
        "state",
        "updateTime",
    }
    return {key: export[key] for key in allowed if key in export}


def summarize_tables(files: list[dict], table_metadata: dict[str, dict] | None = None) -> list[dict]:
    table_metadata = table_metadata or {}
    grouped: dict[str, dict[str, int | str]] = defaultdict(
        lambda: {"file_count": 0, "total_bytes": 0, "schema_file_count": 0, "schema_bytes": 0}
    )
    for item in files:
        name = Path(item["key"]).name
        if name.endswith("-schema-create.sql") or not name.endswith((".parquet", "-schema.sql")):
            continue
        parts = name.split(".")
        if len(parts) < 2:
            continue
        table_name = parts[1].removesuffix("-schema")
        grouped[table_name]["table_name"] = table_name
        if name.endswith(".parquet"):
            grouped[table_name]["file_count"] = int(grouped[table_name]["file_count"]) + 1
            grouped[table_name]["total_bytes"] = int(grouped[table_name]["total_bytes"]) + int(item["size"])
        else:
            grouped[table_name]["schema_file_count"] = int(grouped[table_name]["schema_file_count"]) + 1
            grouped[table_name]["schema_bytes"] = int(grouped[table_name]["schema_bytes"]) + int(item["size"])
    summaries = []
    for table_name, summary in grouped.items():
        summaries.append({**table_metadata.get(table_name, {}), **summary})
    return sorted(summaries, key=lambda item: str(item["table_name"]))


def schema_drift(exported, expected) -> dict[str, list[str]]:
    """Differences between the tables actually exported (parsed from the S3 objects) and
    the tables we have schema metadata/contracts for.

    Surfaced in the manifest so a desync between export and contracts is observable at
    produce time rather than silently dropped (metadata but no files) or deferred to a
    loader hard-fail (files but no metadata, so no column schema).
    """
    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),
    }


def product_status_for(export_state: str) -> str:
    if export_state == "SUCCEEDED":
        return "exported"
    if export_state in {"FAILED", "CANCELED"}:
        return "blocked"
    return "unknown"


def parse_s3_uri(s3_uri: str) -> dict[str, str]:
    parsed = urlparse(s3_uri)
    if parsed.scheme != "s3" or not parsed.netloc:
        raise ValueError(f"not an S3 URI: {s3_uri}")
    return {"bucket": parsed.netloc, "prefix": parsed.path.lstrip("/")}


def run_json(args: list[str], env: dict[str, str] | None = None) -> dict:
    completed = subprocess.run(args, check=True, capture_output=True, text=True, env=env)
    return json.loads(completed.stdout)