#!/usr/bin/env python3
"""
FLUXTRAIT Generation Script — runs ON the RTX PRO 6000 GPU box.
Submits ComfyUI workflows via localhost:8188 API.

Pipeline: FLUXTRAIT checkpoint → 3 skin LoRAs → PuLID @0.65 → KSampler → save + _meta.json
Output: /root/fluxtrait/output/
"""
import json, time, urllib.request, urllib.error, os, sys, traceback, glob

COMFY = "http://localhost:8188"
OUTPUT_DIR = "/home/ubuntu/fluxtrait/output"
REFS_DIR = "/home/ubuntu/fluxtrait/refs"
COMFYUI_INPUT = "/home/ubuntu/ComfyUI/input"  # ComfyUI's LoadImage looks here

# ============================================================
# CONFIG
# ============================================================

CHECKPOINT = "FLUXTRAIT.safetensors"  # fallback: flux1-dev.safetensors
CLIP_L = "clip_l.safetensors"
T5_XXL = "t5xxl_fp8_e4m3fn.safetensors"
PULID_WEIGHT = 0.65
FLUX_GUIDANCE = 4.0  # match proven value from space_generate.py (was 3.5)
STEPS = 24
WIDTH = 832
HEIGHT = 1216

LORA_STACK = [
    {"name": "realistic_skin_texture.safetensors", "strength": 0.6},
    {"name": "skin_no_plastic.safetensors", "strength": 0.6},
    {"name": "detailed_perfection.safetensors", "strength": 0.4},
]

PROMPTS = {
    "studio_portrait": (
        "skin texture style, realism, detailed. aidmarealisticskin. ultra detailed, "
        "detailed skin pore. RAW photo, portrait of a person, "
        "head and shoulders centered in frame, looking directly at camera, "
        "professional studio headshot photograph, soft Rembrandt lighting, "
        "neutral dark gray backdrop. Visible skin pores, natural skin texture "
        "with fine lines, minimal makeup, no foundation, natural imperfections, "
        "micro-texture, subsurface scattering. 85mm lens, f/2.8, shallow depth "
        "of field, professionally color graded, unretouched"
    ),
    "forest_mist": (
        "skin texture style, realism, detailed. aidmarealisticskin. ultra detailed, "
        "detailed skin pore. RAW photo, portrait of a person "
        "standing in a misty ancient forest clearing, head and shoulders visible, "
        "centered in frame, facing the camera, dappled golden sunlight filtering "
        "through the canopy behind them, moss-covered stones in soft bokeh "
        "background. Visible skin pores, natural skin texture, minimal makeup. "
        "85mm portrait lens, shallow depth of field, natural lighting, unretouched"
    ),
    "candlelit_chamber": (
        "skin texture style, realism, detailed. aidmarealisticskin. ultra detailed, "
        "detailed skin pore. RAW photo, portrait of a person "
        "in a dimly lit stone chamber, head and shoulders centered in frame, "
        "facing the camera, warm candlelight casting a golden glow on their face, "
        "rich deep shadows in the background, dramatic chiaroscuro lighting, "
        "medieval fantasy setting. Visible skin pores, natural skin texture. "
        "50mm lens, shallow depth of field, photorealistic, unretouched"
    ),
    "mountain_sunrise": (
        "skin texture style, realism, detailed. aidmarealisticskin. ultra detailed, "
        "detailed skin pore. RAW photo, portrait of a person "
        "on a cliff overlooking a mountain range at sunrise, head and shoulders "
        "centered in frame, facing the camera, wind-swept hair, golden light "
        "bathing their face, dramatic clouds and mountain peaks behind them. "
        "Visible skin pores, natural skin texture, minimal makeup. 85mm portrait "
        "lens, shallow depth of field, golden hour lighting, unretouched"
    ),
}

# ============================================================
# WORKFLOW BUILDER
# ============================================================

def make_workflow(prompt_text, ref_filename, seed):
    """Build FLUXTRAIT + LoRA stack + PuLID workflow JSON."""
    wf = {}

    # Node 1: Checkpoint (model + CLIP — FLUXTRAIT has no VAE, use separate)
    wf["1"] = {
        "inputs": {"ckpt_name": CHECKPOINT},
        "class_type": "CheckpointLoaderSimple",
    }

    # VAELoader — separate Flux VAE (FLUXTRAIT checkpoint doesn't include one)
    wf["vae1"] = {
        "inputs": {"vae_name": "ae.safetensors"},
        "class_type": "VAELoader",
    }

    # Node 2: DualCLIPLoader (clip_l + t5xxl for Flux)
    wf["2"] = {
        "inputs": {
            "clip_name1": CLIP_L,
            "clip_name2": T5_XXL,
            "type": "flux",
        },
        "class_type": "DualCLIPLoader",
    }

    # LoRA chain: checkpoint model (slot 0) + dual clip (slot 0)
    # LoraLoader outputs: slot 0=MODEL, slot 1=CLIP
    prev_model = "1"
    prev_clip = "2"
    prev_clip_slot = 0  # DualCLIPLoader outputs CLIP at slot 0
    for i, lora in enumerate(LORA_STACK):
        node_id = f"lora{i}"
        wf[node_id] = {
            "inputs": {
                "lora_name": lora["name"],
                "strength_model": lora["strength"],
                "strength_clip": lora["strength"],
                "model": [prev_model, 0],
                "clip": [prev_clip, prev_clip_slot],
            },
            "class_type": "LoraLoader",
        }
        prev_model = node_id      # next LoRA reads MODEL from slot 0
        prev_clip = node_id       # next LoRA reads CLIP from slot 1
        prev_clip_slot = 1        # LoraLoader CLIP output is slot 1

    last_lora_node_model = f"lora{len(LORA_STACK) - 1}" if LORA_STACK else "1"
    last_lora_node_clip = f"lora{len(LORA_STACK) - 1}" if LORA_STACK else "2"
    last_lora_clip_slot = 1 if LORA_STACK else 0  # slot 1 from LoRA, slot 0 from DualCLIP

    # Positive + negative text encode
    wf["3"] = {
        "inputs": {"text": prompt_text, "clip": [last_lora_node_clip, last_lora_clip_slot]},
        "class_type": "CLIPTextEncode",
    }
    wf["4"] = {
        "inputs": {"text": "", "clip": [last_lora_node_clip, last_lora_clip_slot]},  # empty neg for Flux
        "class_type": "CLIPTextEncode",
    }

    # Empty latent
    wf["5"] = {
        "inputs": {"width": WIDTH, "height": HEIGHT, "batch_size": 1},
        "class_type": "EmptyLatentImage",
    }

    # Flux guidance
    wf["8"] = {
        "inputs": {"conditioning": ["3", 0], "guidance": FLUX_GUIDANCE},
        "class_type": "FluxGuidance",
    }

    # PuLID nodes
    wf["10"] = {
        "inputs": {"image": ref_filename},
        "class_type": "LoadImage",
    }
    wf["11"] = {
        "inputs": {"pulid_file": "pulid_flux_v0.9.1.safetensors"},
        "class_type": "PulidFluxModelLoader",
    }
    wf["12"] = {
        "inputs": {},
        "class_type": "PulidFluxEvaClipLoader",
    }
    wf["13"] = {
        "inputs": {"provider": "CPU"},
        "class_type": "PulidFluxInsightFaceLoader",
    }
    wf["14"] = {
        "inputs": {
            "model": [last_lora_node_model, 0],   # from LoRA chain or checkpoint
            "pulid_flux": ["11", 0],
            "eva_clip": ["12", 0],
            "face_analysis": ["13", 0],
            "image": ["10", 0],
            "weight": PULID_WEIGHT,
            "start_at": 0.0,
            "end_at": 1.0,
        },
        "class_type": "ApplyPulidFlux",
    }

    # KSampler — uses PuLID-patched model
    wf["6"] = {
        "inputs": {
            "seed": seed, "steps": STEPS, "cfg": 1.0,
            "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0,
            "model": ["14", 0],
            "positive": ["8", 0], "negative": ["4", 0],
            "latent_image": ["5", 0],
        },
        "class_type": "KSampler",
    }

    # VAE decode + save
    wf["7"] = {
        "inputs": {"samples": ["6", 0], "vae": ["vae1", 0]},
        "class_type": "VAEDecode",
    }
    wf["9"] = {
        "inputs": {"filename_prefix": "fluxtrait", "images": ["7", 0]},
        "class_type": "SaveImage",
    }

    return wf


# ============================================================
# COMFYUI API
# ============================================================

def upload_image(filepath):
    """Copy a reference image to ComfyUI's input directory.
    Uses direct file copy (more reliable than upload API, matches the
    proven pattern from space_generate.py)."""
    filename = os.path.basename(filepath)
    dest = os.path.join(COMFYUI_INPUT, filename)
    import shutil
    shutil.copy2(filepath, dest)
    return filename


def submit_and_wait(prompt_text, ref_filename, seed, label):
    """Submit workflow and poll for result."""
    workflow = make_workflow(prompt_text, ref_filename, seed)
    data = json.dumps({"prompt": workflow}).encode()

    req = urllib.request.Request(
        f"{COMFY}/prompt",
        data=data,
        headers={"Content-Type": "application/json"},
    )
    try:
        resp = urllib.request.urlopen(req, timeout=30)
        result = json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body = e.read().decode()[:500]
        return False, f"HTTP {e.code}: {body}"
    except Exception as e:
        return False, f"Submit error: {e}"

    if result.get("node_errors"):
        errs = json.dumps(result["node_errors"])[:500]
        return False, f"Node validation errors: {errs}"

    prompt_id = result.get("prompt_id", "")

    for i in range(180):  # 6 min timeout per image
        time.sleep(2)
        try:
            hist_resp = urllib.request.urlopen(f"{COMFY}/history/{prompt_id}", timeout=10)
            hist = json.loads(hist_resp.read())
        except:
            continue

        if prompt_id not in hist:
            continue

        pd = hist[prompt_id]
        status = pd.get("status", {})
        status_str = status.get("status_str", "")

        if status_str == "error":
            msgs = status.get("messages", [])
            error_detail = ""
            for msg in msgs:
                if isinstance(msg, list) and len(msg) >= 2:
                    if "execution_error" in str(msg[0]):
                        error_detail = json.dumps(msg[1], indent=2)[:2000]
            return False, f"KSampler error: {error_detail or json.dumps(msgs)[:1000]}"

        outputs = pd.get("outputs", {})
        for nid, nout in outputs.items():
            if "images" in nout and nout["images"]:
                return True, nout["images"][0]

        if status_str and status_str != "success":
            return False, f"Unexpected status: {status_str}"

    return False, "Timeout (6 min)"


def download_image(img_info):
    """Download generated image via ComfyUI /view endpoint."""
    view_url = (
        f"{COMFY}/view"
        f"?filename={img_info['filename']}"
        f"&subfolder={img_info.get('subfolder', '')}"
        f"&type={img_info.get('type', 'output')}"
    )
    resp = urllib.request.urlopen(view_url, timeout=30)
    return resp.read()


# ============================================================
# MAIN
# ============================================================

def main():
    import argparse
    parser = argparse.ArgumentParser(description="FLUXTRAIT generation")
    parser.add_argument("--test", action="store_true",
                        help="Run single test image first. Stops after 1 success/failure.")
    parser.add_argument("--full", action="store_true",
                        help="Skip test, run full batch (40 images).")
    args = parser.parse_args()

    # Default: if neither flag, run test first then ask
    test_mode = args.test or not args.full

    # Clean previous run
    import shutil
    if os.path.exists(OUTPUT_DIR):
        shutil.rmtree(OUTPUT_DIR)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Find reference images
    ref_files = sorted(glob.glob(os.path.join(REFS_DIR, "*.png")))
    if not ref_files:
        print(f"❌ No reference images found in {REFS_DIR}")
        print("  Upload reference images before running this script.")
        sys.exit(1)

    print(f"Found {len(ref_files)} reference images")
    for f in ref_files:
        print(f"  - {os.path.basename(f)}")

    # Upload references to ComfyUI input dir
    print("\nUploading reference images to ComfyUI...")
    ref_names = []
    for f in ref_files:
        name = upload_image(f)
        ref_names.append((os.path.basename(f), name))
        print(f"  ✅ {os.path.basename(f)} → {name}")

    total = len(ref_names) * len(PROMPTS)
    done = 0
    failed = 0
    manifest = []
    prompt_keys = list(PROMPTS.keys())

    print(f"\n{'='*60}")
    print(f"FLUXTRAIT Generation: {total} images")
    print(f"  {len(ref_names)} refs × {len(PROMPTS)} prompts")
    print(f"  Checkpoint: {CHECKPOINT}")
    print(f"  LoRAs: {', '.join(l['name'] for l in LORA_STACK)}")
    print(f"  PuLID weight: {PULID_WEIGHT}")
    print(f"  Guidance: {FLUX_GUIDANCE}, Steps: {STEPS}")
    print(f"  Resolution: {WIDTH}×{HEIGHT}")
    print(f"  Output: {OUTPUT_DIR}")
    if test_mode:
        print(f"  ⚠️  TEST MODE: generating 1 image first, then stopping")
    print(f"{'='*60}\n")

    # --- TEST MODE: single image validation ---
    if test_mode:
        ref_orig, ref_comfy = ref_names[0]
        prompt_name = "studio_portrait"
        prompt_text = PROMPTS[prompt_name]
        seed = 42
        label = "fluxtrait_TEST"

        print(f"[TEST] {label} (seed={seed})...", flush=True)
        print(f"  Ref: {ref_orig}", flush=True)
        print(f"  Prompt: {prompt_name}", flush=True)
        print(f"  Checkpoint: {CHECKPOINT}", flush=True)
        print(f"  LoRAs: {LORA_STACK}", flush=True)
        print(f"  PuLID: {PULID_WEIGHT}", flush=True)

        success, result = submit_and_wait(prompt_text, ref_comfy, seed, label)

        if not success:
            print(f"\n❌ TEST FAILED: {result}", flush=True)
            print(f"\nThis error must be fixed before running the full batch.", flush=True)
            with open(os.path.join(OUTPUT_DIR, "TEST_FAILED.txt"), "w") as f:
                f.write(f"Error: {result}\n\n")
                f.write(f"Checkpoint: {CHECKPOINT}\n")
                f.write(f"LoRAs: {json.dumps(LORA_STACK)}\n")
                f.write(f"PuLID: {PULID_WEIGHT}\n")
                f.write(f"Ref: {ref_orig}\n")
            sys.exit(1)

        # Download and save test image
        img_data = download_image(result)
        out_path = os.path.join(OUTPUT_DIR, f"{label}.png")
        with open(out_path, "wb") as f:
            f.write(img_data)

        meta = {
            "positive_prompt": prompt_text,
            "negative_prompt": "",
            "seed": seed, "steps": STEPS, "cfg": 1.0,
            "width": WIDTH, "height": HEIGHT,
            "ip_strength": PULID_WEIGHT, "generator": "rtx6000_comfyui",
            "pulid_weight": PULID_WEIGHT, "reference_image": ref_orig,
            "prompt_type": prompt_name, "checkpoint": CHECKPOINT,
            "loras": [{"name": l["name"], "weight": l["strength"]} for l in LORA_STACK],
        }
        with open(os.path.join(OUTPUT_DIR, f"{label}_meta.json"), "w") as f:
            json.dump(meta, f, indent=2)

        sz_kb = len(img_data) // 1024
        print(f"\n✅ TEST PASSED! Image: {sz_kb}KB", flush=True)
        print(f"   Saved: {out_path}", flush=True)
        print(f"   _meta.json written", flush=True)
        print(f"\n   To run full batch: python3 generate.py --full", flush=True)
        sys.exit(0)

    # --- FULL BATCH MODE ---

    for ref_idx, (ref_orig, ref_comfy) in enumerate(ref_names):
        for prompt_idx, (prompt_name, prompt_text) in enumerate(PROMPTS.items()):
            seed = 10000 * ref_idx + prompt_idx * 1000 + 42
            label = f"fluxtrait_ref{ref_idx}_{prompt_name}"

            print(f"[{done+1}/{total}] {label} (seed={seed})...", flush=True)

            success, result = submit_and_wait(prompt_text, ref_comfy, seed, label)

            if success:
                try:
                    img_data = download_image(result)
                    out_path = os.path.join(OUTPUT_DIR, f"{label}.png")
                    with open(out_path, "wb") as f:
                        f.write(img_data)

                    # Write _meta.json sidecar for the gallery
                    meta = {
                        "positive_prompt": prompt_text,
                        "negative_prompt": "",
                        "seed": seed,
                        "steps": STEPS,
                        "cfg": 1.0,
                        "width": WIDTH,
                        "height": HEIGHT,
                        "ip_strength": PULID_WEIGHT,
                        "generator": "rtx6000_comfyui",
                        "pulid_weight": PULID_WEIGHT,
                        "reference_image": ref_orig,
                        "prompt_type": prompt_name,
                        "checkpoint": CHECKPOINT,
                        "loras": [
                            {"name": l["name"], "weight": l["strength"]}
                            for l in LORA_STACK
                        ],
                        "workflow_json": json.dumps(make_workflow(prompt_text, ref_comfy, seed)),
                    }
                    meta_path = os.path.join(OUTPUT_DIR, f"{label}_meta.json")
                    with open(meta_path, "w") as f:
                        json.dump(meta, f, indent=2)

                    done += 1
                    elapsed_label = f"{len(img_data)//1024}KB"
                    print(f"  ✅ {elapsed_label} + meta → {out_path}", flush=True)
                    manifest.append({"label": label, "ref": ref_orig, "prompt": prompt_name,
                                     "seed": seed, "size_kb": len(img_data)//1024, "status": "ok"})
                except Exception as e:
                    failed += 1
                    print(f"  ❌ Download failed: {e}", flush=True)
                    manifest.append({"label": label, "ref": ref_orig, "prompt": prompt_name,
                                     "seed": seed, "error": str(e)})
            else:
                failed += 1
                print(f"  ❌ FAILED: {result}", flush=True)
                manifest.append({"label": label, "ref": ref_orig, "prompt": prompt_name,
                                 "seed": seed, "error": result})
                # Stop on first KSampler error (usually model/patch issue)
                if "KSampler" in str(result) or "execution_error" in str(result):
                    print("  ⚠️ KSampler error — stopping. Check model files and PuLID patch.")
                    with open(os.path.join(OUTPUT_DIR, "manifest.json"), "w") as f:
                        json.dump(manifest, f, indent=2)
                    sys.exit(1)

            time.sleep(2)  # small delay between gens

    # Write manifest
    with open(os.path.join(OUTPUT_DIR, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)

    # DONE marker for the poller
    with open(os.path.join(OUTPUT_DIR, "DONE"), "w") as f:
        f.write(f"{done} success, {failed} failed")

    print(f"\n{'='*60}")
    print(f"COMPLETE: {done} success, {failed} failed out of {total}")
    print(f"Output: {OUTPUT_DIR}")
    print(f"{'='*60}")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\nFATAL ERROR: {e}", flush=True)
        traceback.print_exc()
        with open(os.path.join(OUTPUT_DIR, "FATAL_ERROR.txt"), "w") as f:
            f.write(f"{e}\n\n{traceback.format_exc()}")
        sys.exit(1)