from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Any


def _load(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _parse_jsonish(value: Any) -> Any:
    if isinstance(value, str):
        s = value.strip()
        if (s.startswith("{") and s.endswith("}")) or (s.startswith("[") and s.endswith("]")):
            try:
                return json.loads(s)
            except Exception:
                return value
    return value


def _extract_harness(payload: dict[str, Any]) -> dict[str, Any]:
    primary = payload.get("hApex")
    fallback = (payload.get("data") or {}).get("_harness_context")
    if isinstance(primary, dict) and isinstance(fallback, dict):
        merged = dict(fallback)
        merged.update(primary)
        return merged
    if isinstance(primary, dict):
        return primary
    if isinstance(fallback, dict):
        return fallback
    return {}


def validate_harness_v4(payload: dict[str, Any]) -> tuple[bool, list[str]]:
    h = _extract_harness(payload)
    required = [
        "fundamental_quality_json",
        "horizon_allocation_json",
        "smart_money_liquidity_json",
        "routing_serving_trace_v2_json",
    ]
    errors: list[str] = []
    for k in required:
        if k not in h:
            errors.append(f"missing key: {k}")

    fq = _parse_jsonish(h.get("fundamental_quality_json", {}))
    if not isinstance(fq, dict) or not isinstance(fq.get("rows"), list):
        errors.append("fundamental_quality_json.rows missing")

    hz = _parse_jsonish(h.get("horizon_allocation_json", {}))
    if not isinstance(hz, dict) or not isinstance(hz.get("bucket_summary"), list):
        errors.append("horizon_allocation_json.bucket_summary missing")

    sml = _parse_jsonish(h.get("smart_money_liquidity_json", {}))
    if not isinstance(sml, dict) or not isinstance(sml.get("rows"), list):
        errors.append("smart_money_liquidity_json.rows missing")

    tr = _parse_jsonish(h.get("routing_serving_trace_v2_json", {}))
    if not isinstance(tr, dict):
        errors.append("routing_serving_trace_v2_json invalid")
    else:
        for k in ("request_route", "json_validation_status"):
            if not tr.get(k):
                errors.append(f"routing_serving_trace_v2_json.{k} missing")
        if tr.get("llm_serving_budget") != 0:
            errors.append("routing_serving_trace_v2_json.llm_serving_budget must be 0")

    return (len(errors) == 0), errors


def main() -> int:
    if len(sys.argv) < 2:
        print("Usage: python tools/apply_engine_upgrade_v4.py <input_json>")
        return 1
    payload = _load(Path(sys.argv[1]))
    ok, errors = validate_harness_v4(payload)
    if not ok:
        print("ENGINE_UPGRADE_V4_FAIL")
        for e in errors:
            print(f"- {e}")
        return 1
    print("ENGINE_UPGRADE_V4_OK")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())