import json
import random
import re
import tempfile
import threading
import zipfile
from pathlib import Path
import gradio as gr
from PIL import Image, ImageOps
from backend import diagnostics, run_one
from catalog import load_catalog

CATALOG = load_catalog()
LOCK = threading.Lock()


def adapters_from_rows(rows, triggers):
    adapters, additions = [], []
    for name, strength in rows:
        strength = float(strength)
        if not name or name == "None" or strength <= 0:
            continue
        if name in CATALOG:
            spec = CATALOG[name]
        else:
            repo, separator, filename = name.partition("::")
            if not separator or "/" not in repo or not filename.endswith((".safetensors", ".gguf")):
                raise gr.Error("Custom LoRA format: owner/repo::folder/file.safetensors")
            spec = {"repo": repo, "file": filename}
        adapters.append((spec, strength))
        if triggers and spec.get("trigger"):
            additions.append(spec["trigger"])
    return adapters, additions


def infer(images, prompt, negative, seed, randomize, cfg, steps, side, triggers,
          *params, progress=gr.Progress()):
    if not images or not prompt.strip():
        raise gr.Error("Upload images and enter an edit prompt.")
    if len(images) > 12:
        raise gr.Error("Maximum 12 images per request.")
    if re.search(r"<lora:", prompt, re.I):
        raise gr.Error("Select LoRAs in the slots rather than embedding LoRA tags in the prompt.")
    adapters, additions = adapters_from_rows(list(zip(params[::2], params[1::2])), triggers)
    effective_prompt = ", ".join([prompt.strip()] + additions)
    folder = Path(tempfile.mkdtemp(prefix="qwen_cpu_"))
    results, failures, manifest = [], [], []
    yield [], None, "Preparing models. First run downloads quantized weights."
    with LOCK:
        for index, uploaded in enumerate(images):
            progress(index / len(images), desc=f"CPU image {index + 1}/{len(images)}")
            current_seed = random.randrange(2**31) if randomize else int(seed)
            record = {"source": Path(uploaded).name, "seed": current_seed,
                      "steps": int(steps), "cfg": float(cfg), "prompt": effective_prompt}
            try:
                with Image.open(uploaded) as opened:
                    image = ImageOps.exif_transpose(opened).convert("RGB")
                    ratio = max(image.size) / min(image.size)
                    if ratio > 4:
                        raise ValueError("Aspect ratio exceeds 4:1")
                    scale = int(side) / max(image.size)
                    width = max(64, round(image.width * scale / 64) * 64)
                    height = max(64, round(image.height * scale / 64) * 64)
                    image = image.resize((width, height), Image.Resampling.LANCZOS)
                    source = folder / f"input_{index:03}.png"
                    image.save(source)
                output = folder / f"edited_{index + 1:03}.png"
                run_one(source, output, effective_prompt, negative, current_seed, int(steps),
                        float(cfg), width, height, adapters, folder,
                        lambda message: progress(index / len(images), desc=message))
                results.append(str(output))
                record.update(status="success", output=output.name, width=width, height=height)
            except Exception as exc:
                record.update(status="failed", error=str(exc))
                failures.append(f"{Path(uploaded).name}: {exc}")
                print(failures[-1], flush=True)
            manifest.append(record)
            summary = f"Completed {index + 1}/{len(images)}; successful {len(results)}; failed {len(failures)}."
            yield list(results), None, summary
    (folder / "manifest.json").write_text(json.dumps(manifest, indent=2))
    bundle = folder / "edited_images.zip"
    with zipfile.ZipFile(bundle, "w", zipfile.ZIP_DEFLATED) as archive:
        for name in results:
            archive.write(name, Path(name).name)
        archive.write(folder / "manifest.json", "manifest.json")
        for log in folder.glob("*.log"):
            archive.write(log, log.name)
    status = f"Finished: {len(results)}/{len(images)} successful."
    if failures:
        status += "\n\n" + "\n\n".join(failures)
    yield results, str(bundle), status


with gr.Blocks(title="Qwen Image Edit CPU Batch", delete_cache=(86400, 86400)) as demo:
    gr.Markdown("# Qwen Image Edit — CPU Batch\nLocal CPU inference; no ZeroGPU or paid inference API. "
                "Quantized base Qwen Edit 2509, not Rapid AIO V4. CPU runs may be very slow. "
                "16 GB RAM is experimental; 32 GB is the recommended initial target.")
    with gr.Row():
        with gr.Column():
            images = gr.File(label="Input images (up to 12)", file_count="multiple", file_types=["image"], type="filepath")
            prompt = gr.Textbox(label="Edit prompt", lines=3)
            negative = gr.Textbox(label="Negative prompt", value="blurry, low quality, watermark")
            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                randomize = gr.Checkbox(label="Random seed per image", value=False)
            cfg = gr.Slider(1, 8, value=2.5, step=0.1, label="CFG")
            steps = gr.Slider(1, 60, value=20, step=1, label="Steps (base model; 4 is not the Rapid preset)")
            side = gr.Slider(256, 1024, value=512, step=64, label="Longest output side")
            triggers = gr.Checkbox(label="Append selected catalogue trigger phrases", value=False)
            slots = []
            with gr.Accordion("Optional LoRA stack — compatibility must be tested", open=False):
                gr.Markdown("Choose an imported entry, or type `owner/repo::path/file.safetensors`. "
                            "Try one LoRA at a time. Four slots do not guarantee four will fit in RAM.")
                for index in range(4):
                    slots += [gr.Dropdown(["None"] + sorted(CATALOG), value="None", allow_custom_value=True,
                                          label=f"LoRA {index + 1}"),
                              gr.Slider(0, 1.5, value=1, step=0.05, label="Strength")]
            run = gr.Button("Process batch on CPU", variant="primary")
        with gr.Column():
            gallery = gr.Gallery(label="Results", columns=2)
            archive = gr.File(label="ZIP: outputs, seeds, metadata and logs")
            status = gr.Textbox(label="Status / errors", lines=8)
    with gr.Accordion("CPU diagnostics", open=False):
        diag = gr.Textbox(label="Backend diagnostics", lines=12)
        check = gr.Button("Check backend (does not load weights)")
        check.click(diagnostics, outputs=diag)
    run.click(infer, [images, prompt, negative, seed, randomize, cfg, steps, side, triggers] + slots,
              [gallery, archive, status], concurrency_limit=1, concurrency_id="cpu-inference")

if __name__ == "__main__":
    demo.queue(max_size=4, default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=7860,
                                                            show_error=True)