"""Portrait-3D — єдиний Gradio ZeroGPU Space (сценарій C, див. ../docs/01).

Чому C замість split (A): щоб списувалась квота ВІДВІДУВАЧА, запит має нести його HF-
ідентичність. HF впорскує її лише у НАТИВНИЙ Gradio-виклик, не в наш сирий fetch. Тож:

  three.js-в'ювер живе в <iframe> (весь WebGL без переписування) → на Generate шле паси
  батьку через postMessage → батько кладе їх у приховані Gradio-inputs і клікає приховану
  кнопку → `@spaces.GPU portrait()` (нативна Gradio-подія = несе HF-токен → per-user квота)
  → результат JSON вертається у прихований textbox → батько postMessage-ить його в iframe.

`pipeline.py`/`passes.py`/`index.html` реюзяться БЕЗ ЗМІН (копіюються сюди `sync.sh`).
"""
from __future__ import annotations

import base64
import html
import io
import json
import os
import time

import gradio as gr

# DEV-режим під Tailscale: якщо задано PORTRAIT_BACKEND (напр. http://pc-gpu:8100), Gradio
# крутиться ЛОКАЛЬНО на Mac, а інференс проксимо на pc-gpu Docker-воркер — без ZeroGPU/torch/
# квоти. Швидка ітерація Gradio-моста без деплою на HF. На HF змінна не задана → ZeroGPU-шлях.
DEV_BACKEND = os.environ.get("PORTRAIT_BACKEND")


def _b64_to_bytes(s: str) -> bytes:
    if not s:
        return b""
    s = s.strip()
    if s.startswith("data:") and "," in s:
        s = s.split(",", 1)[1]
    return base64.b64decode(s)


if DEV_BACKEND:
    # локальний dev: проксі на pc-gpu-воркер (той самий /portrait, той самий pipeline/passes,
    # той самий контракт відповіді) — тому решта коду/моста однакова що тут, що на HF.
    import httpx

    _BOOT = f"DEV proxy → {DEV_BACKEND}"

    def portrait(normal_b64: str, depth_b64: str, beauty_b64: str, params_json: str = "{}",
                 seg_b64: str = "") -> str:
        if not (normal_b64 and depth_b64 and beauty_b64):
            return json.dumps({"error": "normal, depth і beauty — обов'язкові"})
        try:
            _f = [("normal", normal_b64), ("depth", depth_b64), ("beauty", beauty_b64)]
            if seg_b64:  # seg-маска регіонів → воркер нарізає per-part prompt
                _f.append(("seg", seg_b64))
            files = {k: (k + ".png", _b64_to_bytes(v), "image/png") for k, v in _f}
            # connect швидко (мертвий/сплячий pc-gpu → чиста помилка за ~6с, не висне), read довго
            r = httpx.post(f"{DEV_BACKEND}/portrait", files=files, data={"params": params_json},
                           timeout=httpx.Timeout(connect=6.0, read=300.0, write=30.0, pool=6.0))
            return json.dumps(r.json())  # воркер повертає {image,seed,...} → назад як рядок
        except Exception as e:
            return json.dumps({"error": f"dev backend unreachable ({type(e).__name__}) — "
                                        f"wake pc-gpu + run recreate.sh"})

else:
    # HF: нативний ZeroGPU-шлях (моделі в процесі). import spaces ПЕРШИМ — до torch.
    import spaces
    import torch
    from PIL import Image

    import passes
    import pipeline

    DURATION = int(os.environ.get("GPU_DURATION", "60"))
    os.environ.pop("LOW_VRAM", None)  # cpu-offload несумісний із fork-моделлю ZeroGPU

    try:
        pipeline.load()
        pipeline.load_img2img()
        _BOOT = "models loaded at boot"
    except Exception as e:
        _BOOT = f"boot load deferred ({type(e).__name__}: {e})"

    def _png_data_uri(img) -> str:
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()

    @spaces.GPU(duration=DURATION)
    def portrait(normal_b64: str, depth_b64: str, beauty_b64: str, params_json: str = "{}",
                 seg_b64: str = "") -> str:
        """G-буфер-паси + params → JSON {image,seed,timing_ms,meta}|{error}."""
        t0 = time.perf_counter()
        try:
            p = json.loads(params_json or "{}")
        except json.JSONDecodeError as e:
            return json.dumps({"error": f"bad params JSON: {e}"})
        if not (normal_b64 and depth_b64 and beauty_b64):
            return json.dumps({"error": "normal, depth і beauty — обов'язкові"})

        width = int(p.get("width", 832)); height = int(p.get("height", 1216)); size = (width, height)
        cn = p.get("controlnet", {}); flip = p.get("normal_flip", {})
        try:
            normal_img = passes.prep_normal(_b64_to_bytes(normal_b64), flip=flip,
                                            space=p.get("normal_space", "view"), size=size)
            depth_img = passes.prep_depth(_b64_to_bytes(depth_b64),
                                          invert=bool(p.get("depth_invert", False)), size=size)
            beauty_img = passes.prep_beauty(_b64_to_bytes(beauty_b64), size=size)
        except Exception as e:
            return json.dumps({"error": f"bad pass image: {type(e).__name__}"})
        # regional prompting: seg-маска + parts[{color,prompt}] → регіони (паритет із main.py)
        regions = None
        if p.get("use_parts") and seg_b64:
            parts_meta = [pp for pp in p.get("parts", []) if str(pp.get("prompt", "")).strip()]
            if parts_meta:
                regions = passes.build_regions(_b64_to_bytes(seg_b64), parts_meta, size)
        t_prep = time.perf_counter()

        try:
            image, seed = pipeline.generate(
                normal=normal_img, depth=depth_img, beauty=beauty_img,
                prompt=p.get("prompt", ""), negative_prompt=p.get("negative_prompt", ""),
                seed=p.get("seed", -1), steps=int(p.get("steps", 30)),
                cfg_scale=float(p.get("cfg_scale", 6.0)),
                guidance_rescale=float(p.get("guidance_rescale", 0.0)), width=width, height=height,
                normal_scale=float(cn.get("normal_scale", 0.6)), depth_scale=float(cn.get("depth_scale", 0.4)),
                ip_adapter_scale=float(p.get("ip_adapter_scale", 0.6)),
                ip_mode=p.get("ip_mode", "full"), ip_end=float(p.get("ip_end", 1.0)),
                denoise=float(p.get("denoise", 1.0)), regions=regions,
            )
        except torch.cuda.OutOfMemoryError:
            return json.dumps({"error": "GPU out of memory — try smaller size or fewer steps"})
        except Exception as e:
            print(f"generate failed: {type(e).__name__}: {e}", flush=True)
            return json.dumps({"error": f"generate failed: {type(e).__name__}"})
        t_inf = time.perf_counter()

        result = {
            "image": _png_data_uri(image), "seed": seed,
            "timing_ms": {"preprocess": round((t_prep - t0) * 1000),
                          "inference": round((t_inf - t_prep) * 1000),
                          "total": round((t_inf - t0) * 1000)},
            "meta": {"base": pipeline.BASE_ID, "controlnet": pipeline.UNION_ID, "ip_adapter": pipeline.IP_WEIGHT},
        }
        return json.dumps(result)


# --- в'ювер у iframe (srcdoc, self-contained; асети з static-Space через ASSET_BASE) ---
# dev — читаємо ЖИВИЙ ../frontend/index.html (UI-правки без sync); HF — синхронізовану копію.
_HERE = os.path.dirname(__file__)
_VIEWER = next(
    open(p, encoding="utf-8").read()
    for p in (os.path.join(_HERE, "..", "frontend", "index.html"), os.path.join(_HERE, "index.html"))
    if os.path.exists(p)
)
# dev (проксі на pc-gpu) → показати dev-повзунки (Detail/FreeU тощо) у srcdoc-в'ювері,
# де ?dev з URL недоступний. На HF (не dev) — не інжектимо, паблік лишається чистим.
if DEV_BACKEND:
    _VIEWER = _VIEWER.replace("<head>", "<head><script>window.__DEV_KNOBS__=1</script>", 1)
IFRAME = (
    '<iframe id="pv" title="Portrait-3D viewer" '
    'style="width:100%;height:900px;border:0;border-radius:12px;display:block" '
    f'srcdoc="{html.escape(_VIEWER, quote=True)}"></iframe>'
)

# міст батька: приймає паси з iframe → приховані inputs → клік прихованого тригера
BRIDGE = """
<script>
function _postIframe(payload){
  var f=document.querySelector('#pv');
  if(f&&f.contentWindow) f.contentWindow.postMessage({type:'portrait-result', payload:payload}, '*');
}
window.addEventListener('message', function(e){
  var d = e.data || {};
  if (d.type !== 'portrait-generate') return;
  function set(id, val){
    var el = document.querySelector('#'+id+' textarea') || document.querySelector('#'+id+' input');
    if(!el) return;
    var setter = Object.getOwnPropertyDescriptor(el.__proto__, 'value').set;
    setter.call(el, val);
    el.dispatchEvent(new Event('input', {bubbles:true}));
  }
  set('h_normal', d.normal); set('h_depth', d.depth); set('h_beauty', d.beauty); set('h_params', d.params);
  set('h_seg', d.seg || '');
  // gr.Button з elem_id → сам <button> має цей id (не обгортка)
  setTimeout(function(){ var b=document.querySelector('#h_trigger, #h_trigger button'); if(b) b.click(); }, 40);
});
// ZeroGPU-помилки (квота тощо) Gradio показує тостом і НЕ пише в h_result → ловимо тост
// і форвардимо в iframe, щоб в'ювер не зависав на «Generating…».
new MutationObserver(function(muts){
  for (var i=0;i<muts.length;i++) for (var j=0;j<muts[i].addedNodes.length;j++){
    var n=muts[i].addedNodes[j]; if(n.nodeType!==1) continue;
    var t=(n.matches&&n.matches('.toast-body.error'))?n:(n.querySelector&&n.querySelector('.toast-body.error'));
    if(t){ var msg=(t.textContent||'').replace(/\\s*Error\\s*/,'').replace(/^\\s*×\\s*/,'').trim();
      _postIframe({error: msg || 'GPU error'}); }
  }
}).observe(document.body, {childList:true, subtree:true});
</script>
"""

# ⚠️ НЕ visible=False — у Gradio 5 такі компоненти не потрапляють у DOM, і міст їх не знайде.
# Рендеримо нормально, ховаємо CSS-ом → лишаються в DOM, подія тригера все одно спрацьовує.
_HIDE = ("#h_normal,#h_depth,#h_beauty,#h_params,#h_seg,#h_result,#h_trigger"
         "{display:none!important} footer{display:none!important} "
         ".gradio-container{max-width:100%!important;padding:0!important}")
# Gradio 6.0 переніс head/css із Blocks() у launch(). Передаємо туди, куди треба за версією
# (HF пиниться на 5.x, Mac-dev = 6.x) — інакше на 6.x міст (head) мовчки не застосується.
_GR_MAJOR = int(gr.__version__.split(".")[0])
_HEADCSS = {"head": BRIDGE, "css": _HIDE}
with gr.Blocks(title="Portrait-3D", **({} if _GR_MAJOR >= 6 else _HEADCSS)) as demo:
    gr.HTML(IFRAME)
    h_normal = gr.Textbox(elem_id="h_normal")
    h_depth = gr.Textbox(elem_id="h_depth")
    h_beauty = gr.Textbox(elem_id="h_beauty")
    h_params = gr.Textbox(elem_id="h_params")
    h_seg = gr.Textbox(elem_id="h_seg")
    h_result = gr.Textbox(elem_id="h_result")
    h_trigger = gr.Button(elem_id="h_trigger")

    # порядок ВІДПОВІДАЄ data-масиву фронту: [normal, depth, beauty, params, seg]
    h_trigger.click(portrait, [h_normal, h_depth, h_beauty, h_params, h_seg], h_result, api_name="portrait")
    # результат → назад у iframe (js-only подія)
    h_result.change(
        None, h_result, None,
        js="(v)=>{ if(!v) return; var f=document.querySelector('#pv'); "
           "if(f&&f.contentWindow) f.contentWindow.postMessage({type:'portrait-result', payload: JSON.parse(v)}, '*'); }",
    )

demo.queue(max_size=8).launch(**(_HEADCSS if _GR_MAJOR >= 6 else {}))