#!/usr/bin/env python3
from __future__ import annotations

import argparse
import difflib
import json
import re
import shutil
import time
from pathlib import Path

from harness_lib import (
    CODING_TOOLS, SEEDS, Server, answer_matches, args_equal, coding_tool, content, dump,
    execute_tool, final_answer, has, image_message, message, norm,
    score_instruction, score_writing, tests, tool_calls,
)

ROOT = Path("/workspace/kingy-qwen38")


def preflight(a, suite, out):
    s = Server(a.model_id, a.model_path, a.server_bin, out, ctx=8192, reasoning="off", label="preflight").start()
    try:
        run = s.complete([{"role":"system","content":"Follow the exact format."},{"role":"user","content":"Reply with exactly PREFLIGHT_OK"}], seed=SEEDS[0], temperature=0, max_tokens=32)
        dump(out/"preflight.json", {"passed":content(run).strip()=="PREFLIGHT_OK","output":content(run),"raw":run})
    finally:
        dump(out/"preflight-server-summary.json",s.stop())


def capability_preflight(a,suite,out):
    s=Server(a.model_id,a.model_path,a.server_bin,out,ctx=8192,reasoning="on",effort="medium",label="capability-preflight").start()
    try:
        reason=s.complete([{"role":"system","content":"End with FINAL: answer"},{"role":"user","content":"What is 17 + 25?"}],seed=SEEDS[0],temperature=1.0,max_tokens=256)
        case=suite["tools"]["singles"][0]
        tool=s.complete([{"role":"system","content":"Use the supplied tool."},{"role":"user","content":case["prompt"]}],seed=SEEDS[0],temperature=1.0,max_tokens=256,tools=suite["tools"]["tools"])
        calls=tool_calls(message(tool))
        dump(out/"capability-preflight.json",{"reasoning_passed":norm(final_answer(content(reason)))=="42","tool_passed":len(calls)==1 and calls[0]["name"]==case["expected_tool"] and args_equal(calls[0]["args"],case["expected_args"]),"reasoning":reason,"tool":tool,"calls":calls})
    finally:
        dump(out/"capability-preflight-server-summary.json",s.stop())
    if a.mmproj:
        v=Server(a.model_id,a.model_path,a.server_bin,out,ctx=8192,reasoning="off",mmproj=a.mmproj,label="vision-preflight").start()
        try:
            case=suite["vision"][0]
            run=v.complete(image_message(a.root/"fixtures"/case["image"],case["prompt"]),seed=SEEDS[0],temperature=0,max_tokens=128)
            dump(out/"vision-preflight.json",{"passed":has(content(run),case["contains"]),"output":content(run),"raw":run})
        finally:
            dump(out/"vision-preflight-server-summary.json",v.stop())


def nonthinking(a, suite, out):
    s = Server(a.model_id,a.model_path,a.server_bin,out,ctx=8192,reasoning="off",label="nonthinking").start()
    try:
        rows=[]
        for case in suite["instruction"]:
            run=s.complete([{"role":"system","content":"Obey every constraint. Return only the requested answer."},{"role":"user","content":case["prompt"]}],seed=SEEDS[0],temperature=0,max_tokens=256)
            text=content(run); passed,detail=score_instruction(case,text)
            row={"id":case["id"],"passed":passed,"output":text,"detail":detail,"raw":run}
            dump(out/"instruction"/f"{case['id']}.json",row); rows.append(row)
        dump(out/"instruction-summary.json",{"passed":sum(x["passed"] for x in rows),"total":len(rows)})
        rows=[]
        for case in suite["writing"]:
            run=s.complete([{"role":"system","content":"Edit precisely. Preserve facts and all limits."},{"role":"user","content":case["prompt"]}],seed=SEEDS[0],temperature=0,max_tokens=256)
            text=content(run); passed,detail=score_writing(case,text)
            row={"id":case["id"],"passed":passed,"output":text,"detail":detail,"raw":run}
            dump(out/"writing"/f"{case['id']}.json",row); rows.append(row)
        dump(out/"writing-summary.json",{"passed":sum(x["passed"] for x in rows),"total":len(rows)})
        rows=[]
        for case in suite["multilingual"]:
            run=s.complete([{"role":"system","content":"Answer in the requested language and exact format."},{"role":"user","content":case["prompt"]}],seed=SEEDS[0],temperature=0,max_tokens=192)
            text=content(run); passed=has(text,case["contains"])
            row={"id":case["id"],"language":case["language"],"passed":passed,"output":text,"raw":run}
            dump(out/"multilingual"/f"{case['id']}.json",row); rows.append(row)
        dump(out/"multilingual-summary.json",{"passed":sum(x["passed"] for x in rows),"total":len(rows)})
    finally:
        dump(out/"nonthinking-server-summary.json",s.stop())


def reasoning(a,suite,out):
    s=Server(a.model_id,a.model_path,a.server_bin,out,ctx=8192,reasoning="on",effort="medium",label="reasoning").start()
    rows=[]
    try:
        for case in suite["reasoning"]:
            runs=[]
            for seed in SEEDS:
                run=s.complete([{"role":"system","content":"Solve carefully. End with a separate line: FINAL: answer"},{"role":"user","content":case["prompt"]}],seed=seed,temperature=1.0,max_tokens=768)
                text=content(run); final=final_answer(text)
                format_compliant=bool(re.search(r"(?:^|\n)\s*FINAL\s*:\s*.+?\s*$",text,flags=re.I|re.M))
                answer_correct=answer_matches(final,case["answer"])
                passed=format_compliant and answer_correct
                item={"seed":seed,"passed":passed,"answer_correct":answer_correct,"format_compliant":format_compliant,"expected":case["answer"],"final":final,"output":text,"raw":run}
                dump(out/"reasoning"/case["id"]/f"{seed}.json",item); runs.append(item)
            rows.append({"id":case["id"],"passes":sum(x["passed"] for x in runs)})
        dump(out/"reasoning-summary.json",{"passed_runs":sum(x["passes"] for x in rows),"total_runs":len(rows)*len(SEEDS),"problems_any_pass":sum(x["passes"]>0 for x in rows),"problems":len(rows),"rows":rows})
    finally:
        dump(out/"reasoning-server-summary.json",s.stop())


def tools_mode(a,suite,out):
    defs=suite["tools"]["tools"]
    s=Server(a.model_id,a.model_path,a.server_bin,out,ctx=8192,reasoning="on",effort="medium",label="tools").start()
    singles=[]; multis=[]
    try:
        for case in suite["tools"]["singles"]:
            runs=[]
            for seed in SEEDS:
                target=out/"tools-single"/case["id"]/f"{seed}.json"
                if target.exists():
                    runs.append(json.loads(target.read_text()))
                    continue
                try:
                    run=s.complete([{"role":"system","content":"Use the supplied tool. Never invent a result."},{"role":"user","content":case["prompt"]}],seed=seed,temperature=1.0,max_tokens=384,tools=defs)
                    calls=tool_calls(message(run))
                    passed=len(calls)==1 and calls[0]["name"]==case["expected_tool"] and args_equal(calls[0]["args"],case["expected_args"])
                    item={"seed":seed,"passed":passed,"calls":calls,"raw":run}
                except Exception as exc:
                    item={"seed":seed,"passed":False,"calls":[],"error":str(exc)}
                dump(target,item); runs.append(item)
            singles.append({"id":case["id"],"passes":sum(x["passed"] for x in runs)})
        dump(out/"tools-single-summary.json",{"passed_runs":sum(x["passes"] for x in singles),"total_runs":len(singles)*len(SEEDS),"cases_any_pass":sum(x["passes"]>0 for x in singles),"cases":len(singles),"rows":singles})
        for case in suite["tools"]["multis"]:
            runs=[]
            for seed in SEEDS:
                target=out/"tools-multi"/case["id"]/f"{seed}.json"
                if target.exists():
                    runs.append(json.loads(target.read_text()))
                    continue
                messages=[{"role":"system","content":"Complete the task with tools. Never invent results. Give a concise final answer."},{"role":"user","content":case["prompt"]}]
                transcript=[]; expected_i=0; valid=True; final=""; error=None
                for _ in range(6):
                    try:
                        run=s.complete(messages,seed=seed,temperature=1.0,max_tokens=512,tools=defs)
                    except Exception as exc:
                        error=str(exc); valid=False; break
                    msg=message(run); transcript.append(run); calls=tool_calls(msg); messages.append(msg)
                    if not calls:
                        final=msg.get("content") or ""; break
                    for call in calls:
                        if expected_i>=len(case["steps"]):
                            valid=False
                        else:
                            en,ea=case["steps"][expected_i]
                            if call["name"]==en and args_equal(call["args"],ea): expected_i+=1
                            else: valid=False
                        try: result=execute_tool(call["name"],call["args"])
                        except Exception as exc: result={"error":str(exc)}; valid=False
                        messages.append({"role":"tool","tool_call_id":call["id"],"name":call["name"],"content":json.dumps(result)})
                passed=valid and expected_i==len(case["steps"]) and has(final,case["final_contains"])
                item={"seed":seed,"passed":passed,"completed_steps":expected_i,"final":final,"transcript":transcript}
                if error: item["error"]=error
                dump(target,item); runs.append(item)
            multis.append({"id":case["id"],"passes":sum(x["passed"] for x in runs)})
        dump(out/"tools-multi-summary.json",{"passed_runs":sum(x["passes"] for x in multis),"total_runs":len(multis)*len(SEEDS),"cases_any_pass":sum(x["passes"]>0 for x in multis),"cases":len(multis),"rows":multis})
    finally:
        dump(out/"tools-server-summary.json",s.stop())


def coding_mode(a,suite,out):
    s=Server(a.model_id,a.model_path,a.server_bin,out,ctx=32768,reasoning="on",effort="medium",label="coding").start()
    rows=[]
    try:
        for task_file in sorted((a.root/"fixtures"/"coding").glob("*/task.json")):
            task=json.loads(task_file.read_text()); source=task_file.parent; runs=[]
            for seed in SEEDS:
                target=out/"coding"/task["id"]/f"{seed}.json"
                if target.exists():
                    runs.append(json.loads(target.read_text()))
                    continue
                work=out/"coding-workspaces"/task["id"]/str(seed)
                if work.exists(): shutil.rmtree(work)
                shutil.copytree(source,work); before=(work/"app.py").read_text()
                messages=[{"role":"system","content":"Read TASK.md, app.py, test_app.py. Edit only app.py. Run tests until they pass. Do not edit tests."},{"role":"user","content":task["prompt"]}]
                transcript=[]; invalid=0
                for _ in range(10):
                    try:
                        run=s.complete(messages,seed=seed,temperature=1.0,max_tokens=1024,tools=CODING_TOOLS)
                    except Exception as exc:
                        transcript.append({"error":str(exc),"elapsed_seconds":0}); invalid+=1; break
                    msg=message(run); transcript.append(run); calls=tool_calls(msg); messages.append(msg)
                    if not calls: break
                    for call in calls:
                        result=coding_tool(work,call["name"],call["args"])
                        if "error" in result: invalid+=1
                        messages.append({"role":"tool","tool_call_id":call["id"],"name":call["name"],"content":json.dumps(result)})
                test=tests(work); after=(work/"app.py").read_text()
                diff="".join(difflib.unified_diff(before.splitlines(True),after.splitlines(True),fromfile="before/app.py",tofile="after/app.py"))
                item={"seed":seed,"passed":test["exit_code"]==0,"test":test,"invalid_calls":invalid,"diff":diff,"transcript":transcript}
                dump(target,item); runs.append(item)
            rows.append({"id":task["id"],"pass_at_1":runs[0]["passed"],"pass_at_3":any(x["passed"] for x in runs),"passed_runs":sum(x["passed"] for x in runs)})
        dump(out/"coding-summary.json",{"pass_at_1":sum(x["pass_at_1"] for x in rows),"pass_at_3":sum(x["pass_at_3"] for x in rows),"tasks":len(rows),"passed_runs":sum(x["passed_runs"] for x in rows),"total_runs":len(rows)*len(SEEDS),"rows":rows})
    finally:
        dump(out/"coding-server-summary.json",s.stop())


def documents_mode(a,suite,out):
    s=Server(a.model_id,a.model_path,a.server_bin,out,ctx=32768,reasoning="on",effort="medium",label="documents").start()
    rows=[]; token_counts={}
    try:
        for case in suite["documents"]:
            pack=(a.root/"fixtures"/case["pack"]).read_text()
            if case["pack"] not in token_counts: token_counts[case["pack"]]=s.tokenize(pack)
            prompt=f"Answer from SOURCE only. Give answer followed by supporting section ID in square brackets.\n\nSOURCE:\n{pack}\n\nQUESTION: {case['question']}"
            try:
                run=s.complete([{"role":"system","content":"Use only supplied source. If unsupported, say not found."},{"role":"user","content":prompt}],seed=SEEDS[0],temperature=0,max_tokens=256)
                text=content(run); passed=has(text,[case["answer"],case["citation"]])
                row={"id":case["id"],"tier":case["tier"],"passed":passed,"answer":case["answer"],"citation":case["citation"],"pack_tokens":token_counts[case["pack"]],"output":text,"raw":run}
            except Exception as exc:
                row={"id":case["id"],"tier":case["tier"],"passed":False,"answer":case["answer"],"citation":case["citation"],"pack_tokens":token_counts[case["pack"]],"output":"","error":str(exc)}
            dump(out/"documents"/f"{case['id']}.json",row); rows.append(row)
        dump(out/"documents-summary.json",{"passed":sum(x["passed"] for x in rows),"total":len(rows),"token_counts":token_counts})
    finally:
        dump(out/"documents-server-summary.json",s.stop())


def vision_mode(a,suite,out):
    if not a.mmproj: raise ValueError("--mmproj required")
    s=Server(a.model_id,a.model_path,a.server_bin,out,ctx=8192,reasoning="off",mmproj=a.mmproj,label="vision").start()
    rows=[]
    try:
        for case in suite["vision"]:
            msgs=image_message(a.root/"fixtures"/case["image"],case["prompt"]+" Answer concisely.")
            run=s.complete(msgs,seed=SEEDS[0],temperature=0,max_tokens=192)
            text=content(run); passed=has(text,case["contains"])
            row={"id":case["id"],"passed":passed,"expected_contains":case["contains"],"output":text,"raw":run}
            dump(out/"vision"/f"{case['id']}.json",row); rows.append(row)
        dump(out/"vision-summary.json",{"passed":sum(x["passed"] for x in rows),"total":len(rows)})
    finally:
        dump(out/"vision-server-summary.json",s.stop())


def main():
    p=argparse.ArgumentParser()
    p.add_argument("--model-id",required=True); p.add_argument("--model-path",type=Path,required=True)
    p.add_argument("--mmproj",type=Path); p.add_argument("--root",type=Path,default=ROOT)
    p.add_argument("--server-bin",type=Path,default=ROOT/"llama.cpp/build/bin/llama-server")
    p.add_argument("--modes",default="preflight")
    a=p.parse_args(); suite=json.loads((a.root/"fixtures/suite.json").read_text())
    out=a.root/"results"/a.model_id; out.mkdir(parents=True,exist_ok=True)
    dump(out/"run-config.json",{"model_id":a.model_id,"model_path":str(a.model_path),"mmproj":str(a.mmproj) if a.mmproj else None,"server_bin":str(a.server_bin),"modes":a.modes,"suite_version":suite["suite_version"],"seeds":SEEDS})
    dispatch={"preflight":preflight,"capability-preflight":capability_preflight,"nonthinking":nonthinking,"reasoning":reasoning,"tools":tools_mode,"coding":coding_mode,"documents":documents_mode,"vision":vision_mode}
    for mode in [x.strip() for x in a.modes.split(",") if x.strip()]:
        print(json.dumps({"event":"mode_start","model":a.model_id,"mode":mode,"utc":time.strftime("%Y-%m-%dT%H:%M:%SZ",time.gmtime())}),flush=True)
        dispatch[mode](a,suite,out)
        print(json.dumps({"event":"mode_done","model":a.model_id,"mode":mode,"utc":time.strftime("%Y-%m-%dT%H:%M:%SZ",time.gmtime())}),flush=True)


if __name__=="__main__":
    main()
