#!/usr/bin/env python3
import argparse
import json
import os
import subprocess
import tempfile
import time
from pathlib import Path


REQUIRED_KEYS = {
    "understanding_score",
    "satisfaction_score",
    "personal_relevance_score",
    "trust_score",
    "anger_score",
    "policy_support_score",
    "perceived_accountability_score",
    "perceived_fairness_score",
    "perceived_local_fit_score",
    "cost_effectiveness_acceptance_score",
    "preferred_action",
    "response_mode",
    "main_reason",
    "persona_link",
}

SCORE_KEYS = {key for key in REQUIRED_KEYS if key.endswith("_score")}
PREFERRED_ACTIONS = {"continue", "revise", "expand", "reduce", "stop", "unsure"}
RESPONSE_MODES = {
    "user",
    "taxpayer",
    "family_caregiver",
    "worker",
    "local_resident",
    "professional",
    "general_citizen",
    "other",
}


def dump(obj):
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"))


def iter_jsonl(path):
    with Path(path).open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                yield json.loads(line)


def extract_json_object(text):
    text = text.strip()
    if text.startswith("```"):
        text = "\n".join(line for line in text.splitlines() if not line.strip().startswith("```"))
    start = text.find("{")
    end = text.rfind("}")
    if start < 0 or end < start:
        raise ValueError("No JSON object found")
    return json.loads(text[start : end + 1])


def build_messages(row):
    fixed = row["fixed_prompt_prefix"]
    variable = row["variable_input"]
    user_payload = {
        "fixed_prompt_prefix": fixed,
        "VARIABLE INPUT": variable,
        "task": "Read the fixed case pack as the given persona and return JSON matching the output schema.",
    }
    return [
        {"role": "system", "content": fixed["common_instructions"]},
        {"role": "user", "content": json.dumps(user_payload, ensure_ascii=False)},
    ]


def call_openai_compatible(base_url, api_key, model, messages, temperature, max_tokens, timeout):
    url = base_url.rstrip("/") + "/chat/completions"
    body = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "response_format": {"type": "json_object"},
    }
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", delete=False) as body_file:
        json.dump(body, body_file, ensure_ascii=False)
        body_path = body_file.name
    try:
        completed = subprocess.run(
            [
                "curl",
                "-sS",
                "--max-time",
                str(timeout),
                "-H",
                f"Authorization: Bearer {api_key}",
                "-H",
                "Content-Type: application/json",
                "-d",
                f"@{body_path}",
                url,
            ],
            check=False,
            capture_output=True,
            text=True,
        )
    finally:
        try:
            os.unlink(body_path)
        except OSError:
            pass
    if completed.returncode != 0:
        raise RuntimeError(completed.stderr[:500] or f"curl exited {completed.returncode}")
    payload = json.loads(completed.stdout)
    if "error" in payload:
        raise RuntimeError(json.dumps(payload["error"], ensure_ascii=False)[:500])
    message = payload["choices"][0]["message"]
    content = message.get("content") or ""
    # Some reasoning models occasionally place all text in `reasoning` and leave
    # content empty. We do not normally want reasoning text, but using it as a
    # fallback lets retry runs recover JSON if the provider returns it there.
    if not content and isinstance(message.get("reasoning"), str):
        content = message["reasoning"]
    return content, payload.get("usage")


def validate_response(obj):
    problems = []
    if isinstance(obj, list):
        if len(obj) == 1 and isinstance(obj[0], dict):
            obj = obj[0]
        else:
            return ["bad_response_type:list"]
    if not isinstance(obj, dict):
        return [f"bad_response_type:{type(obj).__name__}"]
    missing = sorted(REQUIRED_KEYS - set(obj))
    if missing:
        problems.append(f"missing_keys:{','.join(missing)}")
    extra = sorted(set(obj) - REQUIRED_KEYS)
    if extra:
        problems.append(f"extra_keys:{','.join(extra)}")
    for key in SCORE_KEYS:
        value = obj.get(key)
        if not isinstance(value, int) or not 1 <= value <= 5:
            problems.append(f"bad_score:{key}")
    if obj.get("preferred_action") not in PREFERRED_ACTIONS:
        problems.append("bad_preferred_action")
    if obj.get("response_mode") not in RESPONSE_MODES:
        problems.append("bad_response_mode")
    for key in ("main_reason", "persona_link"):
        value = obj.get(key)
        if not isinstance(value, str) or not value.strip():
            problems.append(f"bad_text:{key}")
        elif len(value) > 180:
            problems.append(f"long_text:{key}")
    return problems


def normalize_response(obj):
    if isinstance(obj, list) and len(obj) == 1 and isinstance(obj[0], dict):
        return obj[0]
    return obj


def make_output_record(row, response, usage, model, problems, simulation_run_id):
    variable = row["variable_input"]
    return {
        "response_id": f"sim-{row['persona_id']}-{row['prompt_hash'][:12]}",
        "simulation_run_id": simulation_run_id,
        "persona_id": row["persona_id"],
        "report_id": row.get("report_id") or row.get("case_pack_id") or "unknown_report",
        "condition_id": row.get("condition_id", "original"),
        "model_id": model,
        "prompt_id": row["prompt_id"],
        "prompt_hash": row["prompt_hash"],
        "cache_key": row["cache_key"],
        "assignment_group": variable.get("assignment_group", ""),
        "outputs": response,
        "usage": usage,
        "validation_problems": problems,
        "created_at_unix": time.time(),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--report", required=True)
    parser.add_argument("--cache-dir", required=True)
    parser.add_argument("--simulation-run-id", default=os.environ.get("SIMULATION_RUN_ID", "simulation_batch_v1"))
    parser.add_argument("--base-url", default=os.environ.get("LLM_BASE_URL", "https://ollama.com/v1"))
    parser.add_argument("--model", default=os.environ.get("LLM_MODEL", "deepseek-v4-flash"))
    parser.add_argument("--api-key-env", default=os.environ.get("LLM_API_KEY_ENV", "OLLAMA_CLOUD_API_KEY"))
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--max-tokens", type=int, default=900)
    parser.add_argument("--retry-max-tokens", type=int, default=4000)
    parser.add_argument("--retries", type=int, default=2)
    parser.add_argument("--flush-every", type=int, default=10)
    parser.add_argument("--timeout", type=int, default=90)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output)
    report_path = Path(args.report)
    cache_dir = Path(args.cache_dir)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)

    rows = list(iter_jsonl(input_path))
    if args.limit:
        rows = rows[: args.limit]

    if args.dry_run:
        dry_path = output_path.with_suffix(".dry_run_prompts.jsonl")
        with dry_path.open("w", encoding="utf-8") as out:
            for row in rows:
                out.write(dump({"persona_id": row["persona_id"], "cache_key": row["cache_key"], "messages": build_messages(row)}) + "\n")
        print(json.dumps({"dry_run_prompts": str(dry_path), "count": len(rows)}, ensure_ascii=False, indent=2))
        return

    api_key = os.environ.get(args.api_key_env)
    if not api_key:
        raise SystemExit(f"Missing API key env: {args.api_key_env}")

    existing = {}
    if output_path.exists():
        for obj in iter_jsonl(output_path):
            existing[obj.get("cache_key")] = obj

    results = list(existing.values())
    report = {
        "input": str(input_path),
        "output": str(output_path),
        "model": args.model,
        "base_url": args.base_url,
        "attempted": 0,
        "cache_hits": 0,
        "written": len(existing),
        "failures": [],
        "validation_problems": {},
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }

    def flush_outputs():
        with output_path.open("w", encoding="utf-8") as out:
            for obj in results:
                out.write(dump(obj) + "\n")
        report["written"] = len(results)
        report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    for row in rows:
        cache_key = row["cache_key"]
        cache_path = cache_dir / (cache_key.replace(":", "_").replace("/", "_") + ".json")
        if cache_key in existing:
            report["cache_hits"] += 1
            continue
        if cache_path.exists():
            cached = json.loads(cache_path.read_text(encoding="utf-8"))
            results.append(cached)
            existing[cache_key] = cached
            report["cache_hits"] += 1
            continue
        report["attempted"] += 1
        last_error = None
        for attempt in range(args.retries + 1):
            try:
                raw_text, usage = call_openai_compatible(
                    args.base_url,
                    api_key,
                    args.model,
                    build_messages(row),
                    args.temperature,
                    args.retry_max_tokens if attempt else args.max_tokens,
                    args.timeout,
                )
                response = normalize_response(extract_json_object(raw_text))
                problems = validate_response(response)
                out = make_output_record(row, response, usage, args.model, problems, args.simulation_run_id)
                if problems:
                    report["validation_problems"][row["persona_id"]] = problems
                if usage:
                    for key in report["usage"]:
                        report["usage"][key] += int(usage.get(key, 0) or 0)
                cache_path.write_text(json.dumps(out, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
                results.append(out)
                existing[cache_key] = out
                last_error = None
                break
            except Exception as exc:
                last_error = f"{type(exc).__name__}: {exc}"
                if attempt < args.retries:
                    time.sleep(2 + attempt * 3)
        if last_error:
            report["failures"].append({"persona_id": row.get("persona_id"), "cache_key": cache_key, "error": last_error})
        if report["attempted"] % max(args.flush_every, 1) == 0:
            flush_outputs()

    flush_outputs()
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()