# Loop
# START:
# -> get a program and hardware context
# -> generate a plan to optimize it
# try the original code
# for each optimization:
#   apply the optimization (LLM)
#   if not (run == sucess && output == correct):
#       retry up to #MAX_FIX_TRIES with feedback
#   else:
#       #feedback the imprv.
# report()
# END:


import json, re, signal, sys
from pathlib import Path
from tools import (
    ask_llm, benchmark_binary, compile_c, compiler_flags,
    hardware_context_text, is_safe_c, run_binary
)
from io_tools import (
    read, write, log, stream, load_plan, save_plan,
    parse_plan, plan_summary, extract_code,
    is_complete, is_sane, outputs_match
)
from prompts import build_plan_prompt, build_apply_prompt, build_fix_prompt


# Paths
ROOT        = Path(__file__).resolve().parent
INPUT       = ROOT / "examples" / "nbody.c"
WORK        = ROOT / "generated"
BASE_SRC    = WORK / "base.c"
BASE_BIN    = WORK / "base"
OPT_SRC     = WORK / "opt.c"
OPT_BIN     = WORK / "opt"
BASE_OUT    = WORK / "base.out"
OPT_OUT     = WORK / "opt.out"
PLAN_FILE   = WORK / "plan.json"
LOG         = WORK / "agent.log"

# Limit the number of retries before giving up an opt.
MAX_FIX_TRIES = 2


# per opt...
# give a prompt -> get the llm raw out -> extract code
# basic sanity check
# complie and run -> compare diff with original -> feedback
def apply_optimization(opt, current_code, baseline_src, hw):
    prompt = build_apply_prompt(current_code, opt, hw, baseline_src)
    raw = ask_llm(prompt, on_chunk=lambda c: stream(c, LOG))
    stream("\n", LOG)

    code = extract_code(raw)

    for attempt in range(MAX_FIX_TRIES + 1):
        complete, reason = is_complete(code)
        if not complete:
            if attempt == MAX_FIX_TRIES:
                return None, "incomplete: " + reason
            log(f"[fix {attempt + 1}] incomplete ({reason})", LOG)
            code = extract_code(
                ask_llm(
                    build_fix_prompt(code, reason, hw, baseline_src),
                    on_chunk=lambda c: stream(c, LOG),
                )
            )
            stream("\n", LOG)
            continue

        sane, msg = is_sane(code)
        if not sane:
            return None, "sanity: " + msg

        safe, msg = is_safe_c(code)
        if not safe:
            return None, "safety: " + msg

        write(OPT_SRC, code)

        ok, cerr = compile_c(str(OPT_SRC), str(OPT_BIN))
        if not ok:
            if attempt == MAX_FIX_TRIES:
                return None, "compile failed: " + cerr[:300]
            log(f"[fix {attempt + 1}] compile error", LOG)
            code = extract_code(
                ask_llm(
                    build_fix_prompt(
                        code,
                        "compile error:\n" + cerr[:600],
                        hw,
                        baseline_src,
                    ),
                    on_chunk=lambda c: stream(c, LOG),
                )
            )
            stream("\n", LOG)
            continue

        correct, diff = outputs_match(
            run_binary, write,
            BASE_BIN, OPT_BIN,
            BASE_OUT, OPT_OUT,
        )

        if not correct:
            if attempt == MAX_FIX_TRIES:
                return None, "wrong output: " + diff
            log(f"[fix {attempt + 1}] wrong output", LOG)
            code = extract_code(
                ask_llm(
                    build_fix_prompt(
                        code,
                        "wrong output:\n" + diff,
                        hw,
                        baseline_src,
                    ),
                    on_chunk=lambda c: stream(c, LOG),
                )
            )
            stream("\n", LOG)
            continue
        return code, ""
    return None, "exceeded fix attempts"

# main loop
def optimize():
    WORK.mkdir(exist_ok=True)
    signal.signal(signal.SIGINT, lambda s, f: sys.exit(130))
    signal.signal(signal.SIGTERM, lambda s, f: sys.exit(130))

    # 1. get hardware details
    hw = hardware_context_text()
    log("[agent] hardware:\n" + hw, LOG)
    log("[agent] flags: " + " ".join(compiler_flags()), LOG)


    # 2. get baseline and copy it complie, run, time
    baseline_src = read(INPUT)
    write(BASE_SRC, baseline_src)
    ok, err = compile_c(str(BASE_SRC), str(BASE_BIN))
    if not ok:
        log("[agent] baseline compile FAILED:\n" + err, LOG)
        return
    base_bench = benchmark_binary(str(BASE_BIN), runs=3)
    base_secs = base_bench["avg_seconds"]
    base_output = base_bench["stdout"].strip()
    write(BASE_OUT, base_bench["stdout"])
    log(f"[agent] baseline: {base_secs:.4f}s", LOG)
    log("[agent] baseline output:\n" + base_output + "\n", LOG)

    # 3. make/read a plan
    plan = load_plan(PLAN_FILE)
    if plan and plan.get("base_secs") == round(base_secs, 6):
        log("[agent] reusing existing plan:", LOG)

    else:
        log("[agent] generating optimization plan...", LOG)
        raw = ask_llm(build_plan_prompt(baseline_src, hw, base_secs), on_chunk=lambda c: stream(c, LOG))
        stream("\n", LOG)
        plan = parse_plan(raw)

        if plan is None:
            log("[agent] failed to parse plan — aborting", LOG)
            return
        for o in plan["optimizations"]:
            o.setdefault("status", "pending")
            o.setdefault("secs", None)
            o.setdefault("note", "")

        plan["base_secs"] = round(base_secs, 6)
        plan["current_code"] = baseline_src
        plan["best_secs"] = base_secs
        save_plan(plan, PLAN_FILE)

    log(plan_summary(plan), LOG)
    log("", LOG)

    # 4. Try each optimization
    current_code = plan.get("current_code") or baseline_src
    best_secs = plan.get("best_secs", base_secs)

    for opt in plan["optimizations"]:
        if opt["status"] in ("ok", "skipped"):
            log(f"[agent] already resolved: [{opt['name']}]", LOG)
            continue
        log(f"\n[agent] applying: [{opt['name']}] {opt['description']}", LOG)
        new_code, fail_reason = apply_optimization(opt, current_code, baseline_src, hw)

        if new_code is None:
            log(f"attempt 1 failed ({fail_reason}) — retrying...", LOG)
            new_code, fail_reason = apply_optimization(
                opt, current_code, baseline_src, hw
            )

        if new_code is None:
            log(f"attempt 2 failed ({fail_reason}) — skipping [{opt['name']}]",LOG)
            opt["status"] = "skipped"
            opt["note"] = fail_reason[:120]
            save_plan(plan, PLAN_FILE)
            continue

        # time
        bench = benchmark_binary(str(OPT_BIN), runs=3)
        secs = bench["avg_seconds"]
        pct = 0 if base_secs <= 0 else -(secs - base_secs) / base_secs * 100

        if pct < -5:
            log(f"[agent] rejected: only {pct:.1f}% improvement (<5%)", LOG)
            opt["status"] = "skipped"
            opt["secs"] = round(secs, 6)
            opt["note"] = f"insufficient improvement {pct:.1f}%"
            save_plan(plan, PLAN_FILE)
            continue

        opt["status"] = "ok"
        opt["secs"] = round(secs, 6)
        opt["note"] = f"{pct:+.1f}% vs baseline"

        current_code = new_code
        plan["current_code"] = current_code

        if secs < best_secs:
            best_secs = secs
            plan["best_secs"] = best_secs

        save_plan(plan, PLAN_FILE)
        log(f"[{opt['name']}] done: {secs:.4f}s ({pct:.1f}%)", LOG)
        log("\n  plan:\n" + plan_summary(plan) + "\n", LOG)

    # 5. report
    pct_final = 0 if base_secs <= 0 else -(best_secs - base_secs) / base_secs * 100
    log("\n" + "=" * 60, LOG)
    log(f"[agent] baseline: {base_secs:.4f}s", LOG)
    log(f"[agent] best:     {best_secs:.4f}s ({pct_final:.1f}%)", LOG)
    log("\nfinal plan:\n" + plan_summary(plan), LOG)
    write(OPT_SRC, current_code)
    write(WORK / "best.c", current_code)
    log("[agent] final code -> generated/best.c", LOG)

if __name__ == "__main__":
    optimize()