"""
FLUX.2 Klein 9B (distilled) + optional uncensored TE + NSFW LoRAs
Lighter than SDXL UnrealVision / Klein-base-50-steps.
"""
import os
import gc
import json
import random
import base64
from io import BytesIO
from pathlib import Path

import gradio as gr
from gradio import Server
from fastapi.responses import HTMLResponse
import numpy as np
import spaces
import torch
from PIL import Image

MAX_SEED = np.iinfo(np.int32).max
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device:", device, "cuda:", torch.cuda.is_available())

from diffusers import Flux2KleinPipeline

dtype = torch.bfloat16

# Distilled 9B = 4 steps, much lighter than base-9B (28–50 steps)
MODEL_ID = os.getenv("KLEIN_MODEL", "black-forest-labs/FLUX.2-klein-9B")
USE_UNCENSORED_TE = os.getenv("KLEIN_UNCENSORED_TE", "1").strip() not in ("0", "false", "no")
UNCENSORED_TE_REPO = os.getenv(
    "KLEIN_TE_REPO",
    "ponpoke/flux2-klein-9b-uncensored-text-encoder",
)

ADAPTER_SPECS = {
    "None": {
        "repo": None,
        "weights": None,
        "adapter_name": None,
        "default_strength": 0.0,
        "hint": "No LoRA",
    },
    "NSFW-Solo": {
        "repo": "diroverflo/FLux_Klein_9B_NSFW",
        "weights": None,  # auto from repo
        "adapter_name": "nsfw-solo",
        "default_strength": 0.75,
        "hint": "Solo NSFW focus (no faces in train set)",
    },
    "Consistency": {
        "repo": "dx8152/Flux2-Klein-9B-Consistency",
        "weights": None,
        "adapter_name": "klein-consistency",
        "default_strength": 0.8,
        "hint": "Identity / edit consistency (not NSFW unlock)",
    },
}

LOADED_ADAPTERS: set = set()
ADAPTER_NAMES = list(ADAPTER_SPECS.keys())

print(f"Loading Flux2KleinPipeline: {MODEL_ID}")
pipe = Flux2KleinPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
)
if torch.cuda.is_available():
    try:
        pipe.enable_model_cpu_offload()
        print("cpu_offload OK")
    except Exception as e:
        print("offload fail:", e)
        pipe = pipe.to(device)
else:
    pipe = pipe.to(device)

# Best-effort uncensored text encoder swap
if USE_UNCENSORED_TE:
    try:
        from transformers import AutoModel, AutoTokenizer

        print(f"[TE] loading uncensored encoder: {UNCENSORED_TE_REPO}")
        tok = AutoTokenizer.from_pretrained(
            UNCENSORED_TE_REPO,
            token=os.getenv("HF_TOKEN") or None,
        )
        te = AutoModel.from_pretrained(
            UNCENSORED_TE_REPO,
            torch_dtype=dtype,
            token=os.getenv("HF_TOKEN") or None,
        )
        # Flux2Klein may expose text_encoder / tokenizer attributes
        if hasattr(pipe, "text_encoder"):
            pipe.text_encoder = te
        if hasattr(pipe, "tokenizer"):
            pipe.tokenizer = tok
        print("[TE] uncensored text encoder attached (best-effort)")
    except Exception as e:
        print(f"[TE] skip uncensored TE ({type(e).__name__}: {e})")
        print("[TE] continuing with official encoder + NSFW LoRAs")

print("Pipeline ready.")


def b64_to_pil_list(b64_json_str):
    if not b64_json_str or str(b64_json_str).strip() in ("", "[]"):
        return []
    try:
        b64_list = json.loads(b64_json_str)
    except Exception:
        return []
    out = []
    for b64_str in b64_list:
        if not b64_str or not isinstance(b64_str, str):
            continue
        try:
            if b64_str.startswith("data:image"):
                _, data = b64_str.split(",", 1)
            else:
                data = b64_str
            out.append(Image.open(BytesIO(base64.b64decode(data))).convert("RGB"))
        except Exception as e:
            print("decode error:", e)
    return out


def pil_to_b64_png(image: Image.Image) -> str:
    buf = BytesIO()
    image.save(buf, format="PNG")
    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"


def update_dimensions(image: Image.Image, max_side: int = 1024):
    if image is None:
        return 1024, 1024
    w, h = image.size
    if w >= h:
        nw = max_side
        nh = int(nw * h / w)
    else:
        nh = max_side
        nw = int(nh * w / h)
    return max(16, (nw // 16) * 16), max(16, (nh // 16) * 16)


app = Server(title="FLUX2-Klein-9B-NSFW")


@app.api(name="check_safety", queue=False)
def check_safety(prompt: str) -> dict:
    return {"status": "ok"}


@app.api(name="generate")
@spaces.GPU(duration=30)
def ping():
    return "gpu-ok"

gr.Interface(ping, None, "text").launch()

    if not prompt or not str(prompt).strip():
        raise gr.Error("Prompt is empty.")

    lora_adapter = (lora_adapter or "NSFW-Solo").strip()
    spec = ADAPTER_SPECS.get(lora_adapter) or ADAPTER_SPECS["None"]
    strength = float(lora_strength if lora_strength is not None else spec["default_strength"])
    strength = max(0.0, min(strength, 1.5))

    if spec["adapter_name"] is None or strength <= 0:
        try:
            pipe.disable_lora()
        except Exception:
            pass
        print("--- LoRA off ---")
    else:
        name = spec["adapter_name"]
        if name not in LOADED_ADAPTERS:
            print(f"--- Loading LoRA {lora_adapter} from {spec['repo']} ---")
            try:
                kwargs = {"adapter_name": name}
                if spec.get("weights"):
                    kwargs["weight_name"] = spec["weights"]
                pipe.load_lora_weights(spec["repo"], **kwargs)
                LOADED_ADAPTERS.add(name)
            except Exception as e:
                raise gr.Error(f"LoRA load failed ({lora_adapter}): {e}")
        try:
            pipe.set_adapters([name], adapter_weights=[strength])
        except Exception:
            pipe.set_adapters([name])
        print(f"--- LoRA {name} @ {strength} ---")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    generator = torch.Generator(
        device="cuda" if torch.cuda.is_available() else "cpu"
    ).manual_seed(seed)

    # Distilled Klein defaults: 4 steps, guidance ~1.0–4.0
    steps = int(max(1, min(int(steps or 4), 28)))
    guidance_scale = float(max(0.0, min(float(guidance_scale or 1.0), 8.0)))

    pil_images = b64_to_pil_list(images_b64_json)
    if pil_images:
        width, height = update_dimensions(pil_images[0])
        processed = [im.resize((width, height), LANCZOS) for im in pil_images]
        image_input = processed if len(processed) > 1 else processed[0]
        print(f"I2I {width}x{height} steps={steps} cfg={guidance_scale}")
    else:
        width, height = 1024, 1024
        image_input = None
        print(f"T2I {width}x{height} steps={steps} cfg={guidance_scale}")

    try:
        kwargs = dict(
            prompt=str(prompt).strip(),
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        )
        if image_input is not None:
            kwargs["image"] = image_input
        result = pipe(**kwargs).images[0]
        return {
            "image": pil_to_b64_png(result),
            "seed": seed,
            "status": "success",
            "width": width,
            "height": height,
            "steps": steps,
            "guidance_scale": guidance_scale,
            "lora": lora_adapter,
            "lora_strength": strength,
        }
    except Exception as e:
        raise gr.Error(f"Inference failed: {type(e).__name__}: {e}")
    finally:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


@app.get("/api/config")
def client_config():
    return {
        "model": MODEL_ID,
        "uncensored_te": USE_UNCENSORED_TE,
        "loras": ADAPTER_NAMES,
        "defaults": {
            "steps": 4,
            "guidance_scale": 1.0,
            "lora": "NSFW-Solo",
            "lora_strength": 0.75,
        },
        "notes": [
            "Distilled Klein 9B — typically 4 steps",
            "TE uncensored is best-effort; LoRA teaches NSFW concepts to DiT",
            "Accept FLUX license on black-forest-labs/FLUX.2-klein-9B",
            "HF_TOKEN if gated models need auth",
        ],
    }


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = Path(__file__).resolve().parent / "index.html"
    if html_path.exists():
        return html_path.read_text(encoding="utf-8")
    return """
<!DOCTYPE html>
<html lang="en"><head>
<meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1"/>
<title>FLUX.2 Klein 9B · NSFW</title>
<style>
:root{--bg:#0b0d12;--card:#141821;--fg:#e8eaed;--muted:#9aa0a6;--a:#a78bfa}
*{box-sizing:border-box}body{margin:0;font-family:Inter,system-ui,sans-serif;background:var(--bg);color:var(--fg)}
.wrap{max-width:920px;margin:0 auto;padding:24px 16px 48px}h1{margin:0 0 4px;font-size:1.45rem}
.sub{color:var(--muted);margin-bottom:16px;font-size:.9rem}
.card{background:var(--card);border:1px solid #252a36;border-radius:14px;padding:14px;margin-bottom:12px}
label{display:block;font-size:.8rem;color:var(--muted);margin-bottom:6px}
textarea,input,select{width:100%;background:#0a0c10;color:var(--fg);border:1px solid #2a2f3a;border-radius:10px;padding:10px}
textarea{min-height:88px}.row{display:grid;grid-template-columns:1fr 1fr;gap:10px}
button{width:100%;border:0;border-radius:12px;padding:12px;font-weight:700;cursor:pointer;
background:linear-gradient(135deg,#7c3aed,#db2777);color:#fff;font-size:1rem}
button:disabled{opacity:.5}#out{max-width:100%;border-radius:12px;margin-top:10px}
#status{white-space:pre-wrap;font-size:.85rem;color:var(--muted)}.hint{font-size:.75rem;color:var(--muted);margin-top:6px}
</style></head><body>
<div class="wrap">
  <h1>FLUX.2 Klein 9B</h1>
  <div class="sub">Distilled · TE uncensored (best-effort) · NSFW LoRA</div>
  <div class="card">
    <label>Image (optional = edit / I2I)</label>
    <input id="file" type="file" accept="image/*"/>
  </div>
  <div class="card">
    <label>Prompt</label>
    <textarea id="prompt">nude woman standing by a window, natural light, photorealistic skin, detailed</textarea>
    <div class="hint">TE uncensored passa o prompt; LoRA ensina anatomia ao DiT.</div>
  </div>
  <div class="card row">
    <div>
      <label>LoRA</label>
      <select id="lora"><option>NSFW-Solo</option><option>Consistency</option><option>None</option></select>
    </div>
    <div>
      <label>LoRA strength</label>
      <input id="str" type="number" value="0.75" min="0" max="1.5" step="0.05"/>
    </div>
    <div>
      <label>Steps (distilled ≈ 4)</label>
      <input id="steps" type="number" value="4" min="1" max="28"/>
    </div>
    <div>
      <label>Guidance</label>
      <input id="cfg" type="number" value="1.0" min="0" max="8" step="0.1"/>
    </div>
  </div>
  <button id="run">Generate</button>
  <div class="card"><div id="status">Ready.</div><img id="out"/></div>
</div>
<script>
const $=id=>document.getElementById(id);
function fileToB64(f){return new Promise((res,rej)=>{const r=new FileReader();r.onload=()=>res(r.result);r.onerror=rej;r.readAsDataURL(f);});}
$("run").onclick=async()=>{
  const btn=$("run");btn.disabled=true;$("status").textContent="Queuing…";$("out").removeAttribute("src");
  try{
    const images=[];const f=$("file").files[0];if(f)images.push(await fileToB64(f));
    const body={data:[JSON.stringify(images),$("prompt").value,$("lora").value,0,true,
      parseFloat($("cfg").value),parseInt($("steps").value,10),parseFloat($("str").value)]};
    const res=await fetch("/gradio_api/call/generate",{method:"POST",headers:{"Content-Type":"application/json"},body:JSON.stringify(body)});
    const ev=await res.json();if(!ev.event_id)throw new Error(JSON.stringify(ev));
    const tr=await fetch(`/gradio_api/call/generate/${ev.event_id}`);
    const text=await tr.text();let payload=null;
    for(const line of text.trim().split("\\n")){if(line.startsWith("data:")){try{payload=JSON.parse(line.slice(5).trim());}catch{}}}
    const out=Array.isArray(payload)?payload[0]:payload;
    if(out&&out.image){$("out").src=out.image;$("status").textContent=`OK · seed=${out.seed} · steps=${out.steps} · lora=${out.lora}@${out.lora_strength}`;}
    else $("status").textContent="Done: "+JSON.stringify(out).slice(0,400);
  }catch(e){$("status").textContent="Error: "+e.message;}finally{btn.disabled=false;}
};
</script></body></html>
"""


if __name__ == "__main__":
    print("Klein 9B distilled + NSFW LoRA — accept BFL license on the model page")
    app.launch(show_error=True, mcp_server=True)