#!/usr/bin/env python3
"""Verify exact lineage from a Tushare scoped manifest to a Qlib provider."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import sys
import tempfile
from typing import Any, Mapping, Optional, Sequence


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from adapters.tushare_local import MANIFEST_NAME, verify_qlib_provider


MATCHED_FIELDS = (
    ("id", "job_id"),
    ("path", "path"),
    ("row_count", "row_count"),
    ("byte_count", "byte_count"),
    ("sha256", "sha256"),
    ("request_sha256", "request_sha256"),
)


class TushareLineageError(RuntimeError):
    """Raised when a scoped manifest and provider do not match exactly."""


def _canonical_json(value: Mapping[str, Any]) -> bytes:
    return json.dumps(
        dict(value),
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode("utf-8")


def _load_object(path: Path, *, label: str) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise TushareLineageError(f"cannot read {label}: {path}") from exc
    if not isinstance(value, dict):
        raise TushareLineageError(f"{label} must be a JSON object")
    return value


def _validate_snapshot(
    snapshot: Mapping[str, Any],
) -> Sequence[Mapping[str, Any]]:
    if snapshot.get("format") != "quant-os.tushare-scoped-snapshot/v2":
        raise TushareLineageError("source snapshot must use scoped format v2")
    manifest_sha = str(snapshot.get("manifest_sha256") or "")
    without_manifest_sha = dict(snapshot)
    without_manifest_sha.pop("manifest_sha256", None)
    computed_manifest_sha = hashlib.sha256(
        _canonical_json(without_manifest_sha)
    ).hexdigest()
    if manifest_sha != computed_manifest_sha:
        raise TushareLineageError("source snapshot manifest SHA-256 mismatch")

    jobs = snapshot.get("selected_jobs")
    if not isinstance(jobs, list) or not jobs:
        raise TushareLineageError("source snapshot has no selected_jobs")
    selection_identity = []
    seen: set[tuple[str, str]] = set()
    for raw_job in jobs:
        if not isinstance(raw_job, Mapping):
            raise TushareLineageError("source snapshot job must be an object")
        key = (
            str(raw_job.get("api_name") or ""),
            str(raw_job.get("partition_key") or ""),
        )
        if not all(key) or key in seen:
            raise TushareLineageError(
                "source snapshot has empty or duplicate job keys"
            )
        seen.add(key)
        path = Path(str(raw_job.get("path") or ""))
        if path.is_absolute() or ".." in path.parts:
            raise TushareLineageError(
                "source snapshot job paths must stay relative to the mirror"
            )
        observed = raw_job.get("observed")
        if raw_job.get("verified") is not True or not isinstance(
            observed,
            Mapping,
        ):
            raise TushareLineageError(
                "every source snapshot job must carry verified observations"
            )
        if (
            observed.get("row_count") != raw_job.get("row_count")
            or observed.get("size_bytes") != raw_job.get("byte_count")
            or observed.get("sha256") != raw_job.get("sha256")
        ):
            raise TushareLineageError(
                "source snapshot observed file evidence is inconsistent"
            )
        selection_identity.append(
            {
                field: value
                for field, value in raw_job.items()
                if field not in {"observed", "verified"}
            }
        )
    computed_selection_sha = hashlib.sha256(
        _canonical_json({"selected_jobs": selection_identity})
    ).hexdigest()
    if computed_selection_sha != snapshot.get("selection_sha256"):
        raise TushareLineageError("source snapshot selection SHA-256 mismatch")

    verification = snapshot.get("verification")
    if (
        not isinstance(verification, Mapping)
        or verification.get("verified_jobs") != len(jobs)
        or verification.get("double_verified_jobs") != len(jobs)
        or verification.get(
            "selected_jobs_still_latest_after_verification"
        )
        is not True
    ):
        raise TushareLineageError(
            "source snapshot lacks closed double-verification evidence"
        )
    return jobs


def _compare_source_jobs(
    snapshot_jobs: Sequence[Mapping[str, Any]],
    provider_jobs: Sequence[Mapping[str, Any]],
) -> None:
    def key(job: Mapping[str, Any]) -> tuple[str, str]:
        return (
            str(job.get("api_name") or ""),
            str(job.get("partition_key") or ""),
        )

    snapshot_by_key = {key(job): job for job in snapshot_jobs}
    provider_by_key = {key(job): job for job in provider_jobs}
    if len(snapshot_by_key) != len(snapshot_jobs):
        raise TushareLineageError("source snapshot has duplicate job keys")
    if len(provider_by_key) != len(provider_jobs):
        raise TushareLineageError("provider manifest has duplicate job keys")
    if set(snapshot_by_key) != set(provider_by_key):
        missing = sorted(set(snapshot_by_key).difference(provider_by_key))
        extra = sorted(set(provider_by_key).difference(snapshot_by_key))
        raise TushareLineageError(
            "source job sets differ: "
            f"missing_from_provider={missing[:5]}, "
            f"extra_in_provider={extra[:5]}"
        )
    mismatches: list[str] = []
    for job_key in sorted(snapshot_by_key):
        source = snapshot_by_key[job_key]
        provider = provider_by_key[job_key]
        for source_field, provider_field in MATCHED_FIELDS:
            if source.get(source_field) != provider.get(provider_field):
                mismatches.append(
                    f"{job_key[0]}:{job_key[1]}:{source_field}"
                )
    if mismatches:
        raise TushareLineageError(
            "source job fields differ: " + ", ".join(mismatches[:12])
        )


def verify_lineage(
    source_manifest: str | Path,
    provider_dir: str | Path,
) -> dict[str, Any]:
    source_path = Path(source_manifest).expanduser().resolve()
    provider = Path(provider_dir).expanduser().resolve()
    snapshot = _load_object(source_path, label="source snapshot")
    snapshot_jobs = _validate_snapshot(snapshot)
    provider_verification = verify_qlib_provider(provider)
    provider_manifest = _load_object(
        provider / MANIFEST_NAME,
        label="provider manifest",
    )
    provider_jobs = provider_manifest.get("source_jobs")
    if not isinstance(provider_jobs, list):
        raise TushareLineageError("provider manifest has no source_jobs")
    _compare_source_jobs(snapshot_jobs, provider_jobs)
    return {
        "artifact_type": "quant_os_tushare_source_provider_lineage",
        "ok": True,
        "snapshot_format": snapshot["format"],
        "snapshot_manifest_sha256": snapshot["manifest_sha256"],
        "snapshot_selection_sha256": snapshot["selection_sha256"],
        "source_job_count": len(snapshot_jobs),
        "matched_fields": [
            source_field for source_field, _ in MATCHED_FIELDS
        ],
        "mismatch_count": 0,
        "provider": provider_verification,
        "tool_source_sha256": hashlib.sha256(
            Path(__file__).resolve().read_bytes()
        ).hexdigest(),
        "gate_credit": [],
        "investment_value_claim": False,
    }


def _write_json(path: Path, payload: Mapping[str, Any]) -> None:
    destination = path.expanduser().resolve()
    if destination.exists() and destination.is_dir():
        raise ValueError("output points to a directory")
    destination.parent.mkdir(parents=True, exist_ok=True)
    provider_bytes = (
        json.dumps(
            dict(payload),
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
            allow_nan=False,
        )
        + "\n"
    ).encode("utf-8")
    descriptor, temporary_name = tempfile.mkstemp(
        prefix=f".{destination.name}.",
        suffix=".tmp",
        dir=destination.parent,
    )
    temporary = Path(temporary_name)
    try:
        with os.fdopen(descriptor, "wb") as handle:
            handle.write(provider_bytes)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, destination)
    finally:
        if temporary.exists():
            temporary.unlink()


def _main(argv: Optional[list[str]] = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-manifest", required=True)
    parser.add_argument("--provider-dir", required=True)
    parser.add_argument("--output-json")
    args = parser.parse_args(argv)
    try:
        result = verify_lineage(
            args.source_manifest,
            args.provider_dir,
        )
        if args.output_json:
            output = Path(args.output_json).expanduser().resolve()
            provider = Path(args.provider_dir).expanduser().resolve()
            if output == provider or provider in output.parents:
                raise ValueError(
                    "output-json must be outside the immutable provider "
                    "directory"
                )
            _write_json(output, result)
    except (OSError, ValueError, TushareLineageError) as exc:
        parser.error(str(exc))
    print(
        json.dumps(
            result,
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
            allow_nan=False,
        )
    )
    return 0


if __name__ == "__main__":
    sys.exit(_main())