"""studio_daemon.py - resident model server for LYRA STUDIO.
Keeps each specialist loaded in VRAM and executes jobs from BASE/jobs.
Writes BASE/results/<id>.json. No time limit."""
import json, os, re, sys, time, traceback
from pathlib import Path
sys.path.insert(0, "/marimo/ai_studio")
import torch
from studio_core import (BASE, CHARACTER, VOICE, SEED, DTYPE_NAME, JOBS, JOBS_DONE, RESULTS, LOGS,
                         build_image_prompt, load_refinements, save_refinements)

DTYPE = getattr(torch, DTYPE_NAME)
MODELS = {}
PAINTER_ID = "Qwen/Qwen-Image"
CRITIC_ID  = "Qwen/Qwen2.5-VL-7B-Instruct"

def log(m):
    line = f"[{time.strftime('%H:%M:%S')}] {m}"
    print(line, flush=True)
    with open(LOGS/"daemon.log", "a") as f: f.write(line + "\n")

def vram():
    free, total = torch.cuda.mem_get_info(); return f"{free/1e9:.0f}/{total/1e9:.0f}GB"

# ---------- specialists ----------
def painter():
    if "painter" not in MODELS:
        from diffusers import QwenImagePipeline
        log("load Qwen-Image (concept artist) ...")
        p = QwenImagePipeline.from_pretrained(PAINTER_ID, torch_dtype=DTYPE)
        p.to("cuda"); p.set_progress_bar_config(disable=True); MODELS["painter"] = p
        log(f"painter ready {vram()}")
    return MODELS["painter"]

def brain():
    if "brain" not in MODELS:
        from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
        log("load Qwen2.5-VL-7B (critic + director brain) ...")
        m = Qwen2_5_VLForConditionalGeneration.from_pretrained(CRITIC_ID, torch_dtype=DTYPE, device_map="cuda")
        pr = AutoProcessor.from_pretrained(CRITIC_ID)
        MODELS["brain"] = (m, pr); log(f"brain ready {vram()}")
    return MODELS["brain"]

# ---------- helpers ----------
def _json_block(s):
    m = re.search(r"\{.*\}", s, re.S)
    if not m: return None
    try: return json.loads(m.group(0))
    except Exception:
        try: return json.loads(m.group(0).replace("'", '"'))
        except Exception: return None

def _gen(model, proc, msgs, images=None, max_new_tokens=900):
    text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    kw = {"text": [text], "return_tensors": "pt"}
    if images: kw["images"] = images
    inputs = proc(**kw).to("cuda")
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    return proc.batch_decode(out[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0].strip()

def _render_one(j, refinements, name):
    pipe = painter()
    prompt = build_image_prompt(j.get("shot",""), j.get("outfit",""), j.get("mood",""), refinements)
    g = torch.Generator(device="cuda").manual_seed(j.get("seed", SEED))
    w, h = j.get("w",1024), j.get("h",1536)
    log(f"render {name} {w}x{h} steps={j.get('steps',30)} {vram()}")
    t0 = time.time()
    img = pipe(prompt=prompt, negative_prompt=CHARACTER["negative"], width=w, height=h,
               num_inference_steps=j.get("steps",30), true_cfg_scale=j.get("cfg",4.0),
               generator=g).images[0]
    out = BASE/"assets/img"/name; img.save(out)
    dt = time.time()-t0; log(f"saved {name} in {dt:.1f}s")
    return {"file": str(out), "prompt": prompt, "seconds": round(dt,1)}

RUBRIC = ("You are the art director of a virtual-idol IP. Be specific and a little ruthless. "
          "Judge the render against the CHARACTER BIBLE. Return STRICT JSON only, keys: "
          "score_identity (0-10), score_beauty (0-10), score_brief (0-10), "
          "problems (list), fixes (list of concrete prompt phrases to ADD to improve the next render), "
          "verdict (one sentence).")

def _critique_one(name, task, max_new_tokens=900):
    from PIL import Image
    model, proc = brain()
    img = Image.open(BASE/"assets/img"/name).convert("RGB")
    user = (f"CHARACTER BIBLE:\n{json.dumps(CHARACTER, ensure_ascii=False, indent=1)}\n\n"
            f"SHOT BRIEF: {task or 'Critique this render of LYRA.'}\n\n{RUBRIC}")
    reply = _gen(model, proc, [{"role":"user","content":[{"type":"image","image":img},
                                                          {"type":"text","text":user}]}],
                 images=[img], max_new_tokens=max_new_tokens)
    return reply, _json_block(reply)

# ---------- jobs ----------
def job_image(j):
    rf = j.get("refinements")
    r = _render_one(j, rf, j["name"])
    return {"files": [r["file"]], "prompt": r["prompt"], "seconds": r["seconds"]}

def job_critique(j):
    reply, parsed = _critique_one(j["image"], j.get("task",""))
    return {"image": j["image"], "critique": reply, "parsed": parsed}

def job_direct(j):
    model, proc = brain()
    log(f"direct: {j['instruction'][:60]}...")
    reply = _gen(model, proc, [{"role":"user","content":[{"type":"text","text":j["instruction"]}]}],
                 max_new_tokens=j.get("max_new_tokens", 900))
    return {"reply": reply}

def job_iterate(j):
    """render -> critique -> append fixes (learned) -> render again."""
    rf = load_refinements()
    history, learned = [], []
    task = j.get("task") or f"Canonical identity shot: {j.get('shot','portrait')}"
    for r in range(int(j.get("rounds", 2))):
        name = j["name"] if r == 0 else j["name"].replace(".png", f"_r{r}.png")
        res = _render_one(j, rf, name)
        reply, parsed = _critique_one(name, task)
        fixes = (parsed or {}).get("fixes", []) if isinstance(parsed, dict) else []
        scores = {k: (parsed or {}).get(k) for k in ("score_identity","score_beauty","score_brief")} if isinstance(parsed, dict) else {}
        history.append({"round": r, "file": res["file"], "prompt": res["prompt"],
                        "critique": reply, "parsed": parsed, "scores": scores})
        log(f"  round {r}: {scores}")
        if fixes: rf = rf + fixes; learned = learned + fixes
    if j.get("learn", True) and learned:
        save_refinements(load_refinements() + learned)
        log(f"learned {len(learned)} refinements -> {len(load_refinements())} total")
    return {"history": history, "final": history[-1]["file"],
            "refinements_file": str(load_refinements())}

def job_evict(j):
    """Free VRAM so another specialist can use the GPU. The studio shares one card."""
    which = j.get("which", "all")
    names = list(MODELS) if which == "all" else ([which] if isinstance(which, str) else list(which))
    for n in names: MODELS.pop(n, None)
    import gc; gc.collect(); torch.cuda.empty_cache()
    log(f"evicted {names} -> {vram()}")
    return {"evicted": names, "vram": vram()}

def job_shutdown(j):
    Path(LOGS/"daemon.stop").write_text("1"); return {"bye": True}

HANDLERS = {"image": job_image, "critique": job_critique, "direct": job_direct,
            "iterate": job_iterate, "evict": job_evict,
            "ping": lambda j: {"pong": True}, "shutdown": job_shutdown}

def main():
    (LOGS/"daemon.stop").unlink(missing_ok=True)
    log(f"daemon up pid={os.getpid()} cuda={torch.cuda.get_device_name(0)} {vram()}")
    while not (LOGS/"daemon.stop").exists():
        found = False
        for f in sorted(JOBS.glob("*.json")):
            found = True; jid = f.stem
            try:
                j = json.loads(f.read_text()); jid = j.get("id", jid)
                h = HANDLERS.get(j.get("type"))
                if h is None:      # not ours (e.g. a tts job) -> leave it for another daemon
                    time.sleep(0.5); continue
                res = h(j)
                res = {"id": jid, "status": "ok", **(res or {})}
            except Exception as e:
                res = {"id": jid, "status": "error", "error": f"{type(e).__name__}: {e}",
                       "trace": traceback.format_exc()[-1500:]}
                log("FAILED " + res["error"])
            tmp = RESULTS/f".{res['id']}.tmp"; tmp.write_text(json.dumps(res, ensure_ascii=False, indent=2))
            tmp.rename(RESULTS/f"{res['id']}.json"); f.rename(JOBS_DONE/f.name)
            log(f"done {res['id']} status={res['status']}")
        if not found: time.sleep(1)

if __name__ == "__main__":
    main()