#!/usr/bin/env python3
"""
experion-loop.py: Hierarchical LLM Loop
  Qwen3.6-27B (Orchestrator)  →  Qwen3-8B (Worker)

Usage:
  python3 experion-loop.py "refactor ExtractTagNames to use compiled regex"
  python3 experion-loop.py "fix all build warnings in src/Core" --max-iter 20
"""

import argparse, json, os, re, subprocess, sys, time, datetime
from openai import OpenAI

QWEN  = dict(base_url="http://localhost:8000/v1",  api_key="none", model="Qwen3.6-27B-FP8")
WORKER_MODEL = dict(base_url="http://localhost:8001/v1", api_key="none", model="Qwen3-8B-FP8")

client = OpenAI(**{k:v for k,v in QWEN.items() if k != "model"}, timeout=300)
worker = OpenAI(**{k:v for k,v in WORKER_MODEL.items() if k != "model"}, timeout=120)

now = lambda: datetime.datetime.now().strftime("%H:%M:%S")

TOOLS = [
    {"type": "function", "function": {
        "name": "read_file", "description": "Read a file",
        "parameters": {"type": "object", "properties": {"filepath": {"type": "string"}}, "required": ["filepath"]}}},
    {"type": "function", "function": {
        "name": "write_file", "description": "Write file (overwrites)",
        "parameters": {"type": "object", "properties": {
            "filepath": {"type": "string"}, "content": {"type": "string"}}, "required": ["filepath", "content"]}}},
    {"type": "function", "function": {
        "name": "edit_file", "description": "Replace exact string in file",
        "parameters": {"type": "object", "properties": {
            "filepath": {"type": "string"}, "old_string": {"type": "string"}, "new_string": {"type": "string"}},
            "required": ["filepath", "old_string", "new_string"]}}},
    {"type": "function", "function": {
        "name": "bash", "description": "Run bash command",
        "parameters": {"type": "object", "properties": {
            "command": {"type": "string"}, "workdir": {"type": "string"}},
            "required": ["command"]}}},
]


def log(msg: str):
    print(f"[{now()}] {msg}", flush=True)


def call_llm(client, model: str, msgs: list, tools=None, max_tok=4096, temp=0.3) -> object:
    kwargs = dict(model=model, messages=msgs, max_tokens=max_tok, temperature=temp)
    if tools:
        kwargs["tools"] = tools
    t0 = time.time()
    r = client.chat.completions.create(**kwargs)
    log(f"LLM {model.split('/')[-1][:12]}… {time.time()-t0:.0f}s  tok={r.usage.total_tokens}")
    msg = r.choices[0].message
    if msg.content is None and msg.tool_calls:
        # tool call response - don't set content
        pass
    elif msg.content is None:
        # reasoning consumed all tokens; retry with higher limit
        log("  ↳ content was None (thinking overflow); retrying with more tokens")
        kwargs["max_tokens"] = max_tok * 2
        t0 = time.time()
        r = client.chat.completions.create(**kwargs)
        log(f"  ↳ retry: {time.time()-t0:.0f}s")
        msg = r.choices[0].message
    return msg


def exec_tool(name: str, args: dict, workdir: str = ".") -> str:
    try:
        if name == "read_file":
            with open(args["filepath"]) as f:
                return f.read()
        elif name == "write_file":
            os.makedirs(os.path.dirname(args["filepath"]) or ".", exist_ok=True)
            with open(args["filepath"], "w") as f:
                f.write(args["content"])
            return f"Written {len(args['content'])}b → {args['filepath']}"
        elif name == "edit_file":
            fp = args["filepath"]; old, new = args["old_string"], args["new_string"]
            with open(fp) as f:
                c = f.read()
            if old not in c:
                return f"Error: 'old_string' not found in {fp}"
            with open(fp, "w") as f:
                f.write(c.replace(old, new, 1))
            return f"Edited {fp}"
        elif name == "bash":
            res = subprocess.run(args["command"], shell=True, capture_output=True, text=True,
                                 cwd=args.get("workdir", workdir), timeout=120)
            out = res.stdout[-4000:] if len(res.stdout) > 4000 else res.stdout
            err = res.stderr[-2000:] if len(res.stderr) > 2000 else res.stderr
            return (out + ("\n[stderr]\n" + err if err else "")).strip() or "(empty)"
        return f"Unknown tool: {name}"
    except Exception as e:
        return f"Error: {e}"


def worker_loop(task: str, workdir: str, max_iter: int = 15) -> dict:
    """Worker loop: repeat until <<DONE>> or max_iter."""
    log("┌─ Worker starting ──────────────────────────────")
    msgs = [{"role": "system", "content": (
        f"You are a coding assistant. Work in: {workdir}\n\nTask:\n{task}\n\n"
        f"Tools: read_file, write_file, edit_file, bash\n"
        f"\n## Protocol (follow strictly):\n"
        f"1. read_file ONCE (first iter only) to understand current code\n"
        f"2. Then use edit_file for ALL targeted changes (NEVER write_file for existing files)\n"
        f"3. Verify with bash (build, test, or check)\n"
        f"4. If verify fails → fix with edit_file → reverify\n"
        f"5. Output <<DONE>> when all criteria met and changes verified\n"
        f"\n## Critical rules:\n"
        f"- read_file: STRICTLY first iteration only. NEVER read the same file twice\n"
        f"  If you need to re-check file content after iteration 1, use 'head' or 'grep' via bash instead\n"
        f"- edit_file: preferred for ALL changes to existing files. NEVER rewrite entire files\n"
        f"- write_file: ONLY for NEW files that don't exist yet\n"
        f"- Make as few tool calls as possible (batch changes)\n"
        f"- Always verify after changes")}]

    for i in range(max_iter):
        log(f"├─ Worker iter {i+1}/{max_iter}")
        try:
            msg = call_llm(worker, WORKER_MODEL["model"], msgs, tools=TOOLS, max_tok=8192, temp=0.2)
        except Exception as e:
            log(f"│  ✗ API error: {e}")
            time.sleep(3)
            continue

        if msg.content and "<<DONE>>" in msg.content:
            log(f"└─ ✓ Worker done ({i+1} iters)")
            return {"status": "done", "iterations": i + 1}

        if msg.tool_calls:
            msgs.append(msg)
            for tc in msg.tool_calls:
                args = json.loads(tc.function.arguments)
                log(f"│  → {tc.function.name} {str(args)[:80]}")
                result = exec_tool(tc.function.name, args, workdir)
                # If this is a read_file, only keep first 1000 chars to avoid context bloat
                truncated = result[:1000] if tc.function.name == "read_file" else result[:5000]
                msgs.append({"role": "tool", "tool_call_id": tc.id, "content": truncated})
        elif msg.content:
            msgs.append(msg)
            msgs.append({"role": "user", "content": "Continue. Make changes and output <<DONE>> when done."})

    log("└─ ✗ Max iterations reached")
    return {"status": "max_iter", "iterations": max_iter}


def orchestrator(goal: str, workdir: str = "."):
    """Qwen plans → Qwen executes → Qwen verifies."""
    log("═══ Qwen: planning ═══")
    msg = call_llm(client, QWEN["model"], [{"role": "system", "content": "You are a planning assistant. Answer concisely in Korean."},
        {"role": "user", "content": (
        f"Goal: {goal}\n\n"
        f"Create a numbered todo list (max 5 items).\n"
        f"Each todo: files to modify, what to change.\n"
        f"Output as markdown numbered list ONLY, no explanation.\n\n"
        f"Project: ExperionCrawler (.NET 8, C#)\n"
        f"  src/Core/, src/Infrastructure/, src/Web/\n"
        f"  Tests: `dotnet test`")}], max_tok=2048)

    plan = msg.content or "(empty plan)"
    print(f"\n{plan}\n")

    todos = [re.sub(r'^\d+[\.\)]\s*', '', l).strip()
             for l in plan.split("\n") if re.match(r'^\d+[\.\)]\s', l)]
    if not todos:
        log("No todos extracted — running full goal as one worker task")
        worker_loop(goal, workdir)
        return

    log(f"═══ Executing {len(todos)} todos ═══")
    prev_results = []
    for idx, todo in enumerate(todos, 1):
        log(f"── Todo {idx}/{len(todos)}: {todo[:80]} ──")
        prev_section = ""
        if prev_results:
            prev_section = "\nPrevious todos completed:\n" + "\n".join(prev_results[-3:]) + "\n"
        msg = call_llm(client, QWEN["model"], [{"role": "user", "content": (
            f"Goal: {goal}\n{prev_section}"
            f"\nTodo ({idx}/{len(todos)}): {todo}\n\n"
            f"Write a concise task description for the worker:\n"
            f"- Files to modify\n- What to change\n- Verification steps\n- Completion criteria\n"
            f"IMPORTANT: Skip verification if the change was already done in a previous todo.")}],
            max_tok=2048)
        task = msg.content or todo
        result = worker_loop(task, workdir)
        log(f"  ↳ {result['status']} ({result['iterations']} iters)")
        prev_results.append(f"Todo {idx}: {todo} — {result['status']} ({result['iterations']} iters)")

    log("═══ Qwen: final verification ═══")
    msg = call_llm(client, QWEN["model"], [{"role": "user", "content": (
        f"Goal: {goal}\n\nAll todos executed. Run final verification:\n"
        f"1. `dotnet build`\n2. `dotnet test`\n3. Report any issues")}],
        tools=[t for t in TOOLS if t["function"]["name"] == "bash"], max_tok=4096)
    if msg.content:
        print(msg.content)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Qwen hierarchical loop (orchestrator + worker)")
    ap.add_argument("goal", nargs="*", help="Goal description")
    ap.add_argument("--dir", "-d", default=".", help="Working directory")
    ap.add_argument("--max-iter", "-m", type=int, default=15, help="Max worker iterations per todo")
    args = ap.parse_args()
    goal = " ".join(args.goal) or "Current dir: review and improve code"
    orchestrator(goal, args.dir)