#!/usr/bin/env python3
"""Validate and compare a controlled first-draft experience-foundation experiment."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys

import jsonschema
import yaml

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.environ.get("CLAUDE_PROJECT_DIR") or os.path.dirname(os.path.dirname(HERE))
SPEC_PATH = os.path.join(ROOT, "org-os", "06-agent-work", "first-draft-experiment-spec.yaml")
EVALUATION_SCHEMA = os.path.join(ROOT, ".claude", "schemas", "first-draft-evaluation.artifact.schema.json")


def _sha(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _load(path):
    with open(path, encoding="utf-8") as handle:
        return yaml.safe_load(handle) or {}


def _resolve(ref, manifest_path):
    if not ref:
        return None
    if os.path.isabs(str(ref)):
        return os.path.normpath(str(ref))
    local = os.path.join(os.path.dirname(os.path.abspath(manifest_path)), str(ref))
    return os.path.normpath(local if os.path.exists(local) else os.path.join(ROOT, str(ref)))


def _live_binding_errors(binding, label, manifest_path):
    path = _resolve((binding or {}).get("ref"), manifest_path)
    if not path or not os.path.isfile(path):
        return [f"{label}: ref 파일 없음"]
    if _sha(path) != (binding or {}).get("sha256"):
        return [f"{label}: live SHA 불일치"]
    return []


def _screenshot_binding_errors(binding, label, manifest_path):
    normalized = {"ref": (binding or {}).get("path"), "sha256": (binding or {}).get("sha256")}
    errors = _live_binding_errors(normalized, label, manifest_path)
    if errors:
        return errors
    path = _resolve(normalized.get("ref"), manifest_path)
    try:
        with open(path, "rb") as handle:
            signature = handle.read(8)
        if signature != b"\x89PNG\r\n\x1a\n" or os.path.getsize(path) <= 1000:
            errors.append(f"{label}: 실제 PNG 시그니처/비자명 크기 증거 필요")
    except OSError as exc:
        errors.append(f"{label}: PNG 검사 실패: {exc}")
    return errors


def _spec():
    return _load(SPEC_PATH).get("first-draft-experiment-spec", {})


def _evaluation(path, manifest_path):
    resolved = _resolve(path, manifest_path)
    document = _load(resolved)
    if isinstance(document.get("payload"), dict):
        document = document["payload"]
    with open(EVALUATION_SCHEMA, encoding="utf-8") as handle:
        schema = json.load(handle)
    errors = sorted(jsonschema.Draft7Validator(schema).iter_errors(document), key=lambda e: list(e.path))
    return resolved, document, [
        f"evaluation {'/'.join(str(value) for value in error.path) or 'payload'}: {error.message}"
        for error in errors
    ]


def validate(manifest_path, require_complete=False):
    root = _load(manifest_path)
    document = root.get("first-draft-experiment") if isinstance(root, dict) else None
    if not isinstance(document, dict):
        return None, ["first-draft-experiment object 없음"]
    errors = []
    status = document.get("status")
    if status not in ("planned", "ready", "running", "completed"):
        errors.append("status는 planned|ready|running|completed")
    if require_complete and status != "completed":
        errors.append("completed experiment 필요")
    control = document.get("control") or {}
    expected_control = _spec().get("control-invariants", {})
    for key, expected in expected_control.items():
        if control.get(key) != expected:
            errors.append(f"control.{key}={expected!r} 고정 필요")
    errors.extend(_live_binding_errors(document.get("request"), "request", manifest_path))
    arms = document.get("arms") or {}
    if set(arms) != {"A", "B"}:
        errors.append("arms는 정확히 A/B")
        return document, errors
    arm_spec = _spec().get("arms", {})
    for arm_id in ("A", "B"):
        arm = arms.get(arm_id) or {}
        expected = arm_spec.get(arm_id) or {}
        if arm.get("treatment") != expected.get("treatment"):
            errors.append(f"arm {arm_id}: treatment 불일치")
        declared = set(arm.get("required-input-kinds") or [])
        missing = set(expected.get("required-input-kinds") or []) - declared
        if missing:
            errors.append(f"arm {arm_id}: required-input-kinds 누락 {sorted(missing)}")
        forbidden = set(arm.get("forbidden-input-kinds") or [])
        missing_forbidden = set(expected.get("forbidden-input-kinds") or []) - forbidden
        if missing_forbidden:
            errors.append(f"arm {arm_id}: forbidden-input-kinds 누락 {sorted(missing_forbidden)}")
        if status in ("ready", "running", "completed"):
            bindings = {item.get("kind"): item for item in arm.get("inputs") or [] if isinstance(item, dict)}
            for kind in declared:
                if kind == "request":
                    continue
                errors.extend(_live_binding_errors(bindings.get(kind), f"arm {arm_id} input {kind}", manifest_path))
        if status == "completed":
            errors.extend(_live_binding_errors(arm.get("output"), f"arm {arm_id} output", manifest_path))
            evaluation_binding = arm.get("evaluation") or {}
            errors.extend(_live_binding_errors(evaluation_binding, f"arm {arm_id} evaluation", manifest_path))
            if not _live_binding_errors(evaluation_binding, "evaluation", manifest_path):
                _path, evaluation, evaluation_errors = _evaluation(evaluation_binding.get("ref"), manifest_path)
                errors.extend(f"arm {arm_id}: {error}" for error in evaluation_errors)
                output = arm.get("output") or {}
                if (evaluation.get("experiment-id") != document.get("experiment-id")
                        or evaluation.get("subject") != document.get("subject")
                        or evaluation.get("arm-id") != arm_id
                        or evaluation.get("model-id") != document.get("model-id")
                        or evaluation.get("request-sha256") != (document.get("request") or {}).get("sha256")
                        or evaluation.get("output-sha256") != output.get("sha256")):
                    errors.append(f"arm {arm_id}: evaluation experiment/model/request/output exact binding 불일치")
                if os.path.abspath(_resolve(evaluation.get("output-ref"), _path) or "") != os.path.abspath(
                        _resolve(output.get("ref"), manifest_path) or ""):
                    errors.append(f"arm {arm_id}: evaluation output-ref가 manifest output-ref와 불일치")
                for viewport in ("desktop", "mobile"):
                    screenshot = (evaluation.get("screenshots") or {}).get(viewport) or {}
                    errors.extend(_screenshot_binding_errors(
                        screenshot, f"arm {arm_id} evaluation screenshot {viewport}", _path))
    if status in ("ready", "running", "completed") and not str(document.get("model-id") or "").strip():
        errors.append("ready 이상은 model-id 고정 필요")
    if (arms.get("A") or {}).get("output", {}).get("ref") == (arms.get("B") or {}).get("output", {}).get("ref") \
            and status == "completed":
        errors.append("A/B output-ref는 서로 달라야 함")
    return document, errors


def plan(manifest_path):
    document, errors = validate(manifest_path)
    if document is None:
        return None, errors
    required = _spec().get("arms", {}).get("B", {}).get("required-input-kinds", [])
    body = {
        "experiment-id": document.get("experiment-id"),
        "status": document.get("status"),
        "model-id": document.get("model-id"),
        "request-sha256": (document.get("request") or {}).get("sha256"),
        "generation-order": ["A", "B"],
        "attempts-per-arm": 1,
        "revision-count-at-capture": 0,
        "arm-B-required-input-kinds": required,
        "ready": not errors and document.get("status") in ("ready", "running", "completed"),
        "validation-errors": errors,
    }
    return body, []


def compare(manifest_path):
    document, errors = validate(manifest_path, require_complete=True)
    if errors:
        return None, errors
    evaluations = {}
    for arm_id in ("A", "B"):
        _path, body, evaluation_errors = _evaluation(document["arms"][arm_id]["evaluation"]["ref"], manifest_path)
        if evaluation_errors:
            return None, evaluation_errors
        evaluations[arm_id] = body["metrics"]
    spec = _spec().get("metrics", {})
    quality = spec.get("score-1-to-5", [])
    lower = spec.get("lower-is-better", [])
    delta = {metric: evaluations["B"][metric] - evaluations["A"][metric] for metric in quality}
    delta.update({metric: evaluations["A"][metric] - evaluations["B"][metric] for metric in lower})
    mean_a = sum(evaluations["A"][metric] for metric in quality) / len(quality)
    mean_b = sum(evaluations["B"][metric] for metric in quality) / len(quality)
    supported = (
        evaluations["B"]["human-preference"] > evaluations["A"]["human-preference"]
        and mean_b > mean_a
        and (evaluations["B"]["revision-count-to-acceptance"] < evaluations["A"]["revision-count-to-acceptance"]
             or evaluations["B"]["tokens-to-acceptance"] < evaluations["A"]["tokens-to-acceptance"])
    )
    return {
        "experiment-id": document.get("experiment-id"),
        "quality-mean": {"A": round(mean_a, 3), "B": round(mean_b, 3)},
        "positive-means-B-better": delta,
        "treatment-supported": supported,
    }, []


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=["validate", "plan", "compare"])
    parser.add_argument("manifest")
    parser.add_argument("--require-complete", action="store_true")
    args = parser.parse_args(argv)
    if args.command == "validate":
        _document, errors = validate(args.manifest, require_complete=args.require_complete)
        result = {"valid": not errors, "errors": errors}
    elif args.command == "plan":
        result, errors = plan(args.manifest)
    else:
        result, errors = compare(args.manifest)
    if errors:
        for error in errors:
            print(f"[first-draft-experiment] ERROR: {error}", file=sys.stderr)
        return 2
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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