#!/usr/bin/env python3
"""Build offline Meta3 workpackets for Qwen/API runs. No network calls."""

from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]


def load_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        data = json.load(handle)
    if not isinstance(data, dict):
        raise ValueError(f"Expected JSON object: {path}")
    return data


def write_json(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def latest_run_dir() -> Path:
    runs_root = ROOT / "runs"
    candidates = sorted(p for p in runs_root.iterdir() if p.is_dir())
    if not candidates:
        raise FileNotFoundError("No run dirs found. Run scripts/meta3_prepare.py first.")
    return candidates[-1]


def load_manifest(run_id: str | None) -> tuple[Path, dict[str, Any]]:
    run_dir = ROOT / "runs" / run_id if run_id else latest_run_dir()
    manifest_path = run_dir / "manifests/run_manifest.json"
    return run_dir, load_json(manifest_path)


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            data = json.loads(line)
            if not isinstance(data, dict):
                raise ValueError(f"Expected JSON object in {path}:{line_no}")
            rows.append(data)
    return rows


def source_manifest_for_run(run_dir: Path) -> dict[str, Any] | None:
    path = run_dir / "manifests/source_snapshot_manifest.json"
    return load_json(path) if path.exists() else None


def source_family(source_path: str) -> str:
    if "/.claude/projects/" in source_path:
        return "claude_projects"
    if "/.hermes/profiles/" in source_path and "/tirith/" in source_path:
        return "hermes_tirith_logs"
    if "/.hermes/profiles/" in source_path and "/.codex/sessions/" in source_path:
        return "hermes_codex_sessions"
    if "/.codex/sessions/" in source_path:
        return "codex_sessions"
    return "unknown"


def stable_digest_key(row: dict[str, Any]) -> tuple[str, str]:
    return (source_family(str(row.get("source_path") or "")), str(row.get("source_path") or ""))


def stratified_sample(rows: list[dict[str, Any]], size: int) -> list[dict[str, Any]]:
    if size <= 0 or size >= len(rows):
        return sorted(rows, key=stable_digest_key)

    buckets: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in sorted(rows, key=stable_digest_key):
        buckets[source_family(str(row.get("source_path") or ""))].append(row)

    selected: list[dict[str, Any]] = []
    family_names = sorted(buckets)
    index = 0
    while len(selected) < size and any(index < len(buckets[name]) for name in family_names):
        for name in family_names:
            bucket = buckets[name]
            if index < len(bucket):
                selected.append(bucket[index])
                if len(selected) >= size:
                    break
        index += 1
    return selected


def compact_digest(row: dict[str, Any], index: int) -> dict[str, Any]:
    source_path = str(row.get("source_path") or "")
    return {
        "digest_index": index,
        "source_path": source_path,
        "source_family": source_family(source_path),
        "created_at": row.get("created_at"),
        "headline": row.get("headline"),
        "what_happened": row.get("what_happened"),
        "tools_agents": row.get("tools_agents", []),
        "outcomes": row.get("outcomes", []),
        "frictions": row.get("frictions", []),
        "patterns": row.get("patterns", []),
        "decisions": row.get("decisions", []),
        "artifacts": row.get("artifacts", []),
        "open_questions": row.get("open_questions", []),
        "evidence": row.get("evidence", []),
        "strategic_relevance": row.get("strategic_relevance", []),
        "confidence": row.get("confidence"),
    }


def chunk_rows(rows: list[dict[str, Any]], shard_size: int) -> list[list[dict[str, Any]]]:
    if shard_size <= 0:
        raise ValueError("--shard-size must be positive")
    return [rows[i : i + shard_size] for i in range(0, len(rows), shard_size)]


def output_contract(question: dict[str, Any], mode: str) -> dict[str, Any]:
    return {
        "format": "json",
        "mode": mode,
        "required_top_level_fields": [
            "question_id",
            "mode",
            "canonical_thesis",
            "claims",
            "counterevidence_status",
            "uncertainty",
            "owner_next_step",
        ],
        "claim_fields": [
            "claim",
            "supporting_digest_indices",
            "supporting_source_paths",
            "counterevidence_or_limit",
            "confidence",
        ],
        "question_answer_contract": question.get("answer_contract", []),
    }


def build_prompt(packet: dict[str, Any]) -> str:
    digest_lines = []
    for digest in packet["digests"]:
        digest_lines.append(
            "\n".join(
                [
                    f"Digest {digest['digest_index']} ({digest['source_family']}, {digest.get('confidence')}):",
                    f"source_path: {digest['source_path']}",
                    f"headline: {digest.get('headline')}",
                    f"what_happened: {digest.get('what_happened')}",
                    f"patterns: {digest.get('patterns')}",
                    f"outcomes: {digest.get('outcomes')}",
                    f"frictions: {digest.get('frictions')}",
                    f"evidence: {digest.get('evidence')}",
                    f"strategic_relevance: {digest.get('strategic_relevance')}",
                ]
            )
        )
    return "\n\n".join(
        [
            "You are evaluating Samuel's agent-session corpus for Meta3.0.",
            "Use only the provided digests as evidence. Do not invent sessions.",
            "Every strong claim needs supporting digest indices and a counterevidence or limit field.",
            "Prefer one canonical owner-useful thesis over many parallel possibilities.",
            f"Mode: {packet['mode']}",
            f"Question {packet['question_id']} - {packet['question_title']}",
            packet["question_prompt"],
            "Return strict JSON matching expected_output_contract.",
            "expected_output_contract:",
            json.dumps(packet["expected_output_contract"], ensure_ascii=False, indent=2),
            "Digests:",
            "\n\n---\n\n".join(digest_lines),
        ]
    )


def build_packets(
    *,
    run_id: str,
    mode: str,
    questions: list[dict[str, Any]],
    digests: list[dict[str, Any]],
    sample_size: int,
    shard_size: int,
    max_packets: int,
) -> list[dict[str, Any]]:
    if mode == "calibration":
        selected = stratified_sample(digests, sample_size)
        shards = [selected]
    elif mode == "claim_ledger":
        selected = sorted(digests, key=stable_digest_key)
        shards = chunk_rows(selected, shard_size)
    else:
        raise ValueError(f"Unsupported mode: {mode}")

    packets: list[dict[str, Any]] = []
    for question in questions:
        for shard_index, shard in enumerate(shards, start=1):
            packet_id = f"{mode}__{question['id']}__shard-{shard_index:04d}"
            compacted = [compact_digest(row, index) for index, row in enumerate(shard, start=1)]
            packet = {
                "schema": "meta3.workpacket.v1",
                "run_id": run_id,
                "packet_id": packet_id,
                "mode": mode,
                "question_id": question["id"],
                "question_title": question.get("title"),
                "question_prompt": question.get("prompt"),
                "question_answer_contract": question.get("answer_contract", []),
                "shard_index": shard_index,
                "shard_count": len(shards),
                "digest_count": len(compacted),
                "digests": compacted,
                "expected_output_contract": output_contract(question, mode),
            }
            packet["prompt"] = build_prompt(packet)
            packets.append(packet)
            if max_packets and len(packets) >= max_packets:
                return packets
    return packets


def build_call_ledger(run_id: str, packets: list[dict[str, Any]]) -> list[dict[str, Any]]:
    created_at = datetime.now(timezone.utc).isoformat()
    return [
        {
            "schema": "meta3.call_ledger_entry.v1",
            "run_id": run_id,
            "packet_id": packet["packet_id"],
            "mode": packet["mode"],
            "question_id": packet["question_id"],
            "status": "planned",
            "attempts": 0,
            "network_calls_made": 0,
            "created_at": created_at,
            "updated_at": created_at,
            "output_path": None,
            "error": None,
        }
        for packet in packets
    ]


def main() -> int:
    parser = argparse.ArgumentParser(description="Build Meta3 offline workpackets without API calls.")
    parser.add_argument("--run-id", default=None, help="Run id under runs/. Defaults to latest run.")
    parser.add_argument("--mode", choices=["calibration", "claim_ledger"], default="calibration")
    parser.add_argument("--sample-size", type=int, default=24)
    parser.add_argument("--shard-size", type=int, default=24)
    parser.add_argument("--max-packets", type=int, default=0, help="0 means no cap.")
    parser.add_argument("--allow-source-drift", action="store_true")
    args = parser.parse_args()

    run_dir, manifest = load_manifest(args.run_id)
    profile = manifest["profile"]
    questions_doc = load_json(Path(manifest["questions_path"]))
    questions = questions_doc.get("questions", [])
    if not isinstance(questions, list) or not questions:
        raise ValueError("No questions found")

    source_manifest = source_manifest_for_run(run_dir)
    if source_manifest:
        digest_path = Path(str(source_manifest["frozen_digest_path"]))
    else:
        digest_path = Path(profile["source_artifacts"]["layer1_digests"])
    digests = load_jsonl(digest_path)
    if not digests:
        raise ValueError(f"No digests found in {digest_path}")

    batch_plan_path = run_dir / "manifests/batch_plan.json"
    if batch_plan_path.exists() and not source_manifest:
        batch_plan = load_json(batch_plan_path)
        source_snapshot = batch_plan.get("source_snapshot", {})
        if isinstance(source_snapshot, dict):
            audit_unique = int(source_snapshot.get("layer1_audit_unique_digest_paths") or 0)
            raw_rows = int(source_snapshot.get("layer1_raw_digest_file_rows_counted") or 0)
            if audit_unique != raw_rows and not args.allow_source_drift:
                raise SystemExit(
                    "Source drift refused: layer1 audit unique digests "
                    f"({audit_unique}) != raw digest rows ({raw_rows}). "
                    "Refresh Meta2 audit or pass --allow-source-drift for an explicit non-final diagnostic run."
                )

    packets = build_packets(
        run_id=manifest["run_id"],
        mode=args.mode,
        questions=questions,
        digests=digests,
        sample_size=args.sample_size,
        shard_size=args.shard_size,
        max_packets=args.max_packets,
    )
    call_ledger = build_call_ledger(manifest["run_id"], packets)

    packets_dir = run_dir / "outputs/workpackets"
    packet_path = packets_dir / f"{args.mode}_workpackets.jsonl"
    call_ledger_path = run_dir / "state" / f"{args.mode}_call_ledger.jsonl"
    packet_manifest_path = packets_dir / f"{args.mode}_packet_manifest.json"
    write_jsonl(packet_path, packets)
    write_jsonl(call_ledger_path, call_ledger)

    source_families: dict[str, int] = defaultdict(int)
    digest_refs_seen: set[str] = set()
    for packet in packets:
        for digest in packet["digests"]:
            source_families[str(digest["source_family"])] += 1
            digest_refs_seen.add(str(digest["source_path"]))

    packet_manifest = {
        "schema": "meta3.workpacket_manifest.v1",
        "run_id": manifest["run_id"],
        "mode": args.mode,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "network_calls_made": 0,
        "source_digest_path": str(digest_path),
        "source_snapshot_manifest_path": str(run_dir / "manifests/source_snapshot_manifest.json") if source_manifest else None,
        "source_snapshot_schema": source_manifest.get("schema") if source_manifest else None,
        "available_digest_rows": len(digests),
        "source_drift_allowed": args.allow_source_drift,
        "packet_count": len(packets),
        "planned_call_count": len(call_ledger),
        "sample_size": args.sample_size,
        "shard_size": args.shard_size,
        "question_count": len(questions),
        "unique_digest_refs_in_packets": len(digest_refs_seen),
        "source_family_mentions": dict(sorted(source_families.items())),
        "packet_path": str(packet_path),
        "call_ledger_path": str(call_ledger_path),
        "expected_response": "strict JSON per packet.expected_output_contract",
        "next_step": "A future network runner may consume this packet file only after explicit approval.",
    }
    if args.mode == "claim_ledger":
        packet_manifest["shard_count_per_question"] = math.ceil(len(digests) / args.shard_size)
    write_json(packet_manifest_path, packet_manifest)

    print(f"packet_manifest={packet_manifest_path}")
    print(f"packet_file={packet_path}")
    print(f"call_ledger={call_ledger_path}")
    print(f"mode={args.mode}")
    print(f"packets={len(packets)}")
    print(f"planned_calls={len(call_ledger)}")
    print(f"unique_digest_refs={len(digest_refs_seen)}")
    print("network_calls_made=0")
    return 0


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