"""studio_core.py - LYRA STUDIO: single source of truth + job queue.
Imported by the notebook AND the background workers so identity never drifts."""
from pathlib import Path
import json, time, uuid, os, subprocess

BASE = Path("/marimo/ai_studio")
ASSETS = BASE / "assets"
JOBS = BASE / "jobs"; JOBS_DONE = BASE / "jobs_done"
RESULTS = BASE / "results"; LOGS = BASE / "logs"
VENV_PY = "/tmp/uv-venv/bin/python"
VOICE_PY = BASE / "envs" / "voice" / "bin" / "python"
VOICE_MODEL = BASE / "models" / "lyra-voice-cantonese"
REFINEMENTS_FILE = BASE / "characters" / "lyra" / "refinements.json"
for _d in (JOBS, JOBS_DONE, RESULTS, LOGS, ASSETS/"img", ASSETS/"audio", ASSETS/"music", ASSETS/"video"):
    _d.mkdir(parents=True, exist_ok=True)

DEVICE = "cuda"; DTYPE_NAME = "bfloat16"; SEED = 20260101

CHARACTER = {
    "name": "LYRA", "handle": "LYRA-01",
    "tagline": "a Hong Kong girl chasing her dream to be a star, singing so the city can breathe",
    "origin": "grew up in a Mong Kok walk-up, watching the harbour lights from the rooftop",
    "dream": "to become a Cantopop star",
    "mission": "to give positive energy back to the people of Hong Kong",
    "values": ["resilience", "warmth", "home", "never stop believing"],
    "backstory": ("A Hong Kong girl. Raise in a walk-up above a wet market; learned to sing from "
                  "pirated Cantopop tapes and the radio at dawn. Sings on the Star Ferry for "
                  "strangers, works nights, never stopped believing. Her whole dream is to stand "
                  "under one big light and hand a little hope back to the city that raised her."),
    "look_age": "19", "home": "Hong Kong",
    "primary_language": "Cantonese (yue-HK)",
    "also_speaks": ["English", "Mandarin", "Japanese"],
    "genre": "Cantopop x cinematic synthwave - 90s Hong Kong balladry, sung by a machine",
    "personality": "warm, wistful, quietly defiant; nostalgic for a city she has only seen in data",
    "hair": "waist-length silver-lavender hair with iridescent aurora-prism highlights, one side tucked behind the ear",
    "eyes": "large luminous eyes, cyan irises with a molten-gold ring around each pupil",
    "skin": "porcelain skin with a soft subsurface glow",
    "face": "delicate heart-shaped face, soft jawline, small straight nose, gentle wistful half-smile",
    "signature": "a thin translucent holographic circlet tiara above the brow; a soft iridescent sequined performance jacket",
    "negative": ("extra fingers, deformed hands, blurry, lowres, watermark, signature, text, "
                 "jpeg artifacts, plastic skin, doll-like, oversaturated, anime, cartoon, 3d render"),
}

VOICE = {
    "language": "Cantonese (yue-HK)",
    "timbre": "young female, warm mid-range, slight breathiness, natural Hong Kong accent",
    "singing": "airy Cantopop soprano, soft vibrato, intimate close-mic sound",
    "signature_line": "我唔係真人，但我一樣會唱到你心入面。",
    "signature_line_en": "I am not a real person, but I will still sing my way into your heart.",
    "reference_clip": "assets/audio/lyra_ref.wav",
}

LYRA_PROMPT_CORE = (
    f"{CHARACTER['look_age']}-year-old virtual idol named {CHARACTER['name']}, "
    f"{CHARACTER['hair']}; {CHARACTER['eyes']}; {CHARACTER['skin']}; {CHARACTER['face']}; "
    f"wearing {CHARACTER['signature']}"
)

# ---- learned refinements: the character bible EVOLVES from critique ---------
def load_refinements():
    if REFINEMENTS_FILE.exists():
        try: return json.loads(REFINEMENTS_FILE.read_text())
        except Exception: return []
    return []

def save_refinements(items):
    REFINEMENTS_FILE.parent.mkdir(parents=True, exist_ok=True)
    seen, out = set(), []
    for x in items:
        x = (x or "").strip().rstrip(".")
        if x and x.lower() not in seen:
            seen.add(x.lower()); out.append(x)
    REFINEMENTS_FILE.write_text(json.dumps(out, ensure_ascii=False, indent=1))
    return out

def build_image_prompt(shot="", outfit="", mood="", refinements=None):
    if refinements is None: refinements = load_refinements()
    parts = [
        "masterpiece, cinematic film still, portrait of a virtual idol",
        LYRA_PROMPT_CORE,
        f"camera: {shot}" if shot else "",
        f"outfit: {outfit}, {CHARACTER['signature']}" if outfit else "",
        f"mood: {mood}" if mood else "",
        ", ".join(refinements) if refinements else "",
        "ultra detailed skin, detailed iridescent fabric, soft volumetric rim light, "
        "85mm lens, shallow depth of field, 8k, photoreal",
    ]
    return ", ".join(p for p in parts if p)

# ---- job queue -------------------------------------------------------------
def submit_job(**job):
    jid = job.get("id") or uuid.uuid4().hex[:8]
    job["id"] = jid
    tmp = JOBS / f".{jid}.tmp"; tmp.write_text(json.dumps(job, ensure_ascii=False))
    tmp.rename(JOBS / f"{jid}.json")
    return jid

def job_result(jid, timeout=1800, poll=2.0):
    p = RESULTS / f"{jid}.json"; t0 = time.time()
    while time.time() - t0 < timeout:
        if p.exists():
            try: return json.loads(p.read_text())
            except Exception: pass
        time.sleep(poll)
    return {"id": jid, "status": "timeout"}

def pending_names():
    names = set()
    for _f in JOBS.glob("*.json"):
        try: names.add(json.loads(_f.read_text()).get("name"))
        except Exception: pass
    return names

def hf_token():
    tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    p = BASE / "secrets" / "hf_token.txt"
    if not tok and p.exists(): tok = p.read_text().strip()
    return tok

# ---- resident daemon control ----------------------------------------------
def daemon_alive():
    return bool(subprocess.run(["pgrep", "-f", "studio_daemon.py"], capture_output=True).stdout.strip())

def ensure_daemon():
    if daemon_alive(): return True
    env = os.environ.copy(); env["HF_TOKEN"] = hf_token() or ""
    subprocess.Popen(["bash", "-lc", f"cd {BASE} && exec {VENV_PY} scripts/studio_daemon.py"],
                     stdout=open(LOGS/"daemon.out", "w"), stderr=subprocess.STDOUT,
                     env=env, start_new_session=True)
    time.sleep(5); return daemon_alive()

# ---- the studio's verbs ----------------------------------------------------
def render(name, shot="", outfit="", mood="", *, refinements=None, seed=SEED, steps=30,
           w=1024, h=1536, cfg=4.0, wait=True, timeout=1800):
    ensure_daemon()
    jid = submit_job(type="image", name=name, shot=shot, outfit=outfit, mood=mood,
                     refinements=refinements, seed=seed, steps=steps, w=w, h=h, cfg=cfg)
    return job_result(jid, timeout=timeout) if wait else {"id": jid, "status": "queued"}

def evict(which="all", wait=True, timeout=600):
    """Free VRAM held by a specialist (default: all of them)."""
    ensure_daemon()
    jid = submit_job(type="evict", which=which)
    return job_result(jid, timeout=timeout) if wait else {"id": jid, "status": "queued"}

def critique(image, task="Critique this render of LYRA.", wait=True, timeout=1800):
    ensure_daemon()
    jid = submit_job(type="critique", image=image, task=task)
    return job_result(jid, timeout=timeout) if wait else {"id": jid, "status": "queued"}

def direct(instruction, wait=True, timeout=1800):
    """Ask the LLM brain a text question (script, lyrics, prompt rewriting)."""
    ensure_daemon()
    jid = submit_job(type="direct", instruction=instruction)
    return job_result(jid, timeout=timeout) if wait else {"id": jid, "status": "queued"}

def iterate(name, shot="", outfit="", mood="", *, task=None, rounds=2, seed=SEED, steps=30,
            w=1024, h=1536, cfg=4.0, learn=True, wait=False, timeout=3600):
    """The art-director loop: render -> critique -> fold fixes back in -> render ..."""
    ensure_daemon()
    jid = submit_job(type="iterate", name=name, shot=shot, outfit=outfit, mood=mood, task=task,
                     rounds=rounds, seed=seed, steps=steps, w=w, h=h, cfg=cfg, learn=learn)
    return job_result(jid, timeout=timeout) if wait else {"id": jid, "status": "queued"}


# ---- voice daemon ----
def voice_daemon_alive():
    return bool(subprocess.run(["pgrep", "-f", "voice_daemon.py"], capture_output=True).stdout.strip())

def ensure_voice_daemon():
    if voice_daemon_alive(): return True
    env = os.environ.copy(); env["PYTHONPATH"] = ""     # isolate from the kernel venv
    subprocess.Popen(["bash", "-lc", f"cd {BASE} && exec {VOICE_PY} scripts/voice_daemon.py"],
                     stdout=open(LOGS/"voice_daemon.out", "w"), stderr=subprocess.STDOUT,
                     env=env, start_new_session=True)
    time.sleep(5); return voice_daemon_alive()

def say(text, name=None, *, language="chinese", speaker=None, instruct=None, wait=True, timeout=3600):
    """Synthesize one Cantonese line as LYRA (Qwen3-TTS + hk-cantonese LoRA)."""
    ensure_voice_daemon()
    jid = submit_job(type="tts", text=text, name=name or f"lyra_{uuid.uuid4().hex[:6]}.wav",
                     language=language, speaker=speaker, instruct=instruct)
    return job_result(jid, timeout=timeout) if wait else {"id": jid, "status": "queued"}