import os
import sys
import multiprocessing

if __name__ == "__main__":
    multiprocessing.freeze_support()

EARLY_IS_FROZEN = getattr(sys, "frozen", False)
EARLY_BASE_DIR = os.path.dirname(sys.executable) if EARLY_IS_FROZEN else os.path.dirname(os.path.abspath(__file__))
EARLY_APP_SUPPORT_DIR = os.path.join(
    os.path.expanduser("~"),
    "Library",
    "Application Support",
    "3d-model-generator",
)
EARLY_HF_HOME = os.path.join(EARLY_APP_SUPPORT_DIR if EARLY_IS_FROZEN else EARLY_BASE_DIR, ".hf_cache")
os.environ.setdefault("HF_HOME", EARLY_HF_HOME)
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.path.join(EARLY_HF_HOME, "hub"))

import json
import re
import requests
import threading
import time
import base64
import io
import webview
import random
import subprocess
import concurrent.futures
import shutil
from PIL import Image, PngImagePlugin
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch
import replicate

# ------------- Configuration -------------

OLLAMA_URL = "http://localhost:11434/api/generate"
MODEL = "mistral:latest"  # or "mistral-small3.1:24b"

# Path to your diffusers-converted model
MODEL_PATH = "/Volumes/SD/ML-Models/diffusers/dreamshaper_8_diffusers"

IS_FROZEN = getattr(sys, "frozen", False)
BASE_DIR = os.path.dirname(sys.executable) if IS_FROZEN else os.path.dirname(os.path.abspath(__file__))
BUNDLE_DIR = getattr(sys, "_MEIPASS", BASE_DIR)
APP_SUPPORT_DIR = EARLY_APP_SUPPORT_DIR
SETTINGS_PATH = os.path.join(APP_SUPPORT_DIR, "settings.json")
# ────────────────────────────────────────────────────────────────────────────


def resource_path(*parts: str) -> str:
    return os.path.join(BUNDLE_DIR, *parts)


def prepare_web_dir() -> str:
    if not IS_FROZEN:
        return os.path.join(BASE_DIR, "web")

    src = resource_path("web")
    dst = os.path.join(APP_SUPPORT_DIR, "web")
    os.makedirs(dst, exist_ok=True)

    for root, dirs, files in os.walk(src):
        dirs[:] = [d for d in dirs if d != "output"]
        rel = os.path.relpath(root, src)
        out_root = dst if rel == "." else os.path.join(dst, rel)
        os.makedirs(out_root, exist_ok=True)
        for fname in files:
            shutil.copy2(os.path.join(root, fname), os.path.join(out_root, fname))

    return dst


WEB_DIR = prepare_web_dir()
WEB_INDEX = os.path.join(WEB_DIR, "index.html")
OUTPUT_DIR = os.path.join(WEB_DIR, "output")
APP_ICON_PATH = resource_path("icon.png")
os.makedirs(OUTPUT_DIR, exist_ok=True)

IMAGE_BACKEND_LOCAL = "local"
IMAGE_BACKEND_REPLICATE = "replicate"
DEFAULT_IMAGE_BACKEND = IMAGE_BACKEND_LOCAL
DEFAULT_REPLICATE_IMAGE_MODEL = "google/imagen-4-fast"
REPLICATE_IMAGE_TIMEOUT = 10 * 60
REPLICATE_IMAGE_POLL_INTERVAL = 2.0

REPLICATE_IMAGE_MODELS = {
    "black-forest-labs/flux-schnell": {
        "label": "FLUX.1 schnell",
        "price": "$3 / 1000 output images (~$0.003/image)",
        "aspect_ratios": ["1:1", "16:9", "21:9", "3:2", "2:3", "4:5", "5:4", "3:4", "4:3", "9:16", "9:21"],
    },
    "google/imagen-4-fast": {
        "label": "Google Imagen 4 Fast",
        "price": "$0.02 / output image",
        "aspect_ratios": ["1:1", "9:16", "16:9", "3:4", "4:3"],
    },
    "bytedance/seedream-4": {
        "label": "ByteDance Seedream 4",
        "price": "$0.03 / output image",
        "aspect_ratios": ["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9"],
    },
}


def load_app_settings() -> dict:
    try:
        with open(SETTINGS_PATH, "r", encoding="utf-8") as f:
            settings = json.load(f)
            return settings if isinstance(settings, dict) else {}
    except FileNotFoundError:
        return {}
    except Exception:
        return {}


def save_app_settings(settings: dict) -> None:
    os.makedirs(APP_SUPPORT_DIR, exist_ok=True)
    with open(SETTINGS_PATH, "w", encoding="utf-8") as f:
        json.dump(settings, f, indent=2)

# ------------- Prompt Templates -------------

META_PROMPT = """
You are an expert at writing concise, detailed Stable Diffusion prompts for 3D rendered objects, as seen in professional game development.

When I give you an object name (which may already contain style or material cues, e.g. "low poly farmer" or "pixel art sword"), follow these rules:

1. Parse the object name:
- If the object name includes a style, technique, or specific look (e.g., low poly, pixel art, voxel, cartoon, realistic, hand painted, clay, wireframe, etc.), make that style the focus of the prompt.
- Do NOT add conflicting style keywords. (E.g. do NOT add high poly if the input is low poly.)
- If the object name is generic (like farmer or spaceship), you may suggest professional, detailed, high-quality, realistic or stylized game art (choose one style, do NOT mix).
- Add 2–3 relevant visual or material details to the object, based on typical game art conventions.

2. Medium & Style:
- If the object name includes a medium, engine, or tool (Unreal, Octane, Unity, voxel, etc.), use it and do NOT add others.
- Otherwise, you may choose (but do NOT mix) one or two: concept art, 3D render, game asset, professional game designer, digital sculpture, etc.

3. Presentation:
- State that the object is fully visible, not cropped, and centered from a neutral angle.
- No background: isolated on pure white or transparent background.

4. Lighting & Quality:
- Only add studio lighting, sharp focus, clean silhouette, ultra detailed, 8K, etc. if compatible with the style.
- Avoid depth of field, blur, or effects that obscure details, unless specifically requested.

5. Negative Prompt:
- Always add: Negative prompt: blurry, lowres, bad anatomy, distorted proportions, background, scenery, environment, artifacting, watermark, text, cropped, partial view, depth of field, out of frame, blur, vignette

Output Format:
- Line 1: All positive prompt keywords (subject, details, medium, style, lighting, quality, presentation, background).
- Line 2: Negative prompt: then the negative keywords above.

Instructions:
- If the input object already specifies a style or technique, your prompt must reflect that and avoid all contradictory terms.
- Always keep the style consistent.
- Avoid unnecessary repetition.
- Never add stylistic terms that contradict or dilute the user’s intention.
- The positive prompt must not contain more than 300 characters. Make it just fit 300 characters, not too long, not too short but perfect for a Stable Diffusion prompt.
- If the prompt is short, come up with creative ideas on how to make the character or object more intricate and interesting, suitable for game design. Add creative details to what the presented things might look like!
- Use positive descriptions of what should be visible in the positive prompt, don't use negations there.
- Most importantly, the prompts must make sure in some way that the object or character is FULLY visible (not cropped) on a blank background.

Now, generate a Stable Diffusion prompt for the object: {{OBJECT}}
"""

SECOND_PROMPT_TEMPLATE = """
You are a prompt-to-JSON converter for image generation tasks.

Given a Stable Diffusion prompt, extract the positive and negative prompts.
If the prompt is longer than 300 characters, summarize it so that it's less than 300 characters long.
If you need to shorten your prompt, always remove less relevant visual details first.
Never remove the presentation instructions (e.g., fully visible, isolated, centered, white background); these must always remain at the end of the positive prompt.
Remove words that are followed by a semicolon (i.e. "Presentation:" or "Lighting & Quality:") - these things have no business in a Stable Diffusion prompt.

Also, analyze the subject:
- If it is a single character, an upright (bipedal) creature, full object, a person, or anything similar, set "dimensions" to [512,768] (vertical). (In that case of a bipedal creature / human, you might also add "A-Pose" to the prompt!)
- If it is a landscape, wide object, multi-character group, or scene, or a creature that isn't tall (four-legged creatures for example), set "dimensions" to [768,512] (horizontal).
- If it is something square or best seen as a square (e.g., shield, logo, face, emblem, single centered item or a box or round-shaped creature), set "dimensions" to [768,768] (square).

Return a JSON object with these fields:
{
  "positive prompt": "...",
  "negative prompt": "...",
  "dimensions": [W,H]
}

Respond only with the JSON.
Here is the Stable Diffusion prompt:

{PROMPT}
"""

# ------------- Helper Functions -------------

def stream_ollama(prompt: str):
    """
    Stream tokens from Ollama.
    Yields the cumulative response string after each new chunk arrives.
    """
    payload = {"model": MODEL, "prompt": prompt, "stream": True}
    with requests.post(OLLAMA_URL, json=payload, stream=True, timeout=600) as r:
        r.raise_for_status()
        full_response = ""
        for line in r.iter_lines():
            if line:
                try:
                    chunk = json.loads(line.decode("utf-8"))["response"]
                except Exception:
                    chunk = line.decode("utf-8")
                full_response += chunk
                yield full_response
        yield full_response.strip()

def get_json_from_llm(prompt: str):
    """
    Stream the JSON conversion call to Ollama (no LangChain).
    Yields the cumulative response string after each chunk arrives.
    """
    payload = {"model": MODEL, "prompt": prompt, "stream": True}
    with requests.post(OLLAMA_URL, json=payload, stream=True, timeout=600) as r:
        r.raise_for_status()
        full_response = ""
        for line in r.iter_lines():
            if line:
                try:
                    chunk = json.loads(line.decode("utf-8"))["response"]
                except Exception:
                    chunk = line.decode("utf-8")
                full_response += chunk
                yield full_response
        yield full_response.strip()

def sanitize_filename(s: str) -> str:
    s = s.strip().lower().replace(" ", "_")
    return re.sub(r'[^a-z0-9_]+', '', s)

def get_next_available_filename(base: str, ext: str = ".png") -> str:
    i = 1
    while True:
        fname = f"{base}-{i}{ext}"
        # → hier steht noch fpath = f"output/{fname}"
        fpath = os.path.join(OUTPUT_DIR, fname)
        if not os.path.exists(fpath):
            return fpath
        i += 1

def image_to_base64(path: str) -> str:
    """
    Read an image from disk and return a data: URI string.
    """
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    return f"data:image/png;base64,{b64}"


def normalize_image_backend(value: str | None) -> str:
    value = (value or DEFAULT_IMAGE_BACKEND).strip().lower()
    if value in {IMAGE_BACKEND_LOCAL, IMAGE_BACKEND_REPLICATE}:
        return value
    return DEFAULT_IMAGE_BACKEND


def normalize_replicate_image_model(value: str | None) -> str:
    value = (value or DEFAULT_REPLICATE_IMAGE_MODEL).strip()
    if value in REPLICATE_IMAGE_MODELS:
        return value
    return DEFAULT_REPLICATE_IMAGE_MODEL


def nearest_aspect_ratio(width: int, height: int, allowed: list[str]) -> str:
    if width <= 0 or height <= 0:
        return "1:1" if "1:1" in allowed else allowed[0]

    target = width / height

    def ratio_value(value: str) -> float:
        left, right = value.split(":", 1)
        return float(left) / float(right)

    return min(allowed, key=lambda option: abs(ratio_value(option) - target))


def compose_replicate_prompt(prompt_json: dict) -> str:
    positive = str(prompt_json.get("positive prompt", "")).strip()
    negative = str(prompt_json.get("negative prompt", "")).strip()
    if not negative:
        return positive
    return f"{positive}\nAvoid: {negative}"


def extract_replicate_output_urls(output) -> list[str]:
    if output is None:
        return []
    if hasattr(output, "url"):
        return [str(output.url)]
    if isinstance(output, str):
        return [output]
    if isinstance(output, (list, tuple)):
        urls = []
        for item in output:
            urls.extend(extract_replicate_output_urls(item))
        return urls
    if isinstance(output, dict):
        urls = []
        for item in output.values():
            urls.extend(extract_replicate_output_urls(item))
        return urls
    return []


def load_image_from_url(url: str) -> Image.Image:
    if url.startswith("data:"):
        _, encoded = url.split(",", 1)
        data = base64.b64decode(encoded)
    else:
        response = requests.get(url, timeout=REPLICATE_IMAGE_TIMEOUT)
        response.raise_for_status()
        data = response.content

    img = Image.open(io.BytesIO(data))
    if img.mode in {"RGBA", "LA"} or (img.mode == "P" and "transparency" in img.info):
        return img.convert("RGBA")
    return img.convert("RGB")


def save_generated_image(img: Image.Image, prompt_json: dict, object_name: str, idx: int, batch_count: int) -> str:
    fname = get_next_available_filename(f"{sanitize_filename(object_name)}-{idx}")
    metadata = dict(prompt_json, created=int(time.time()), batch_count=batch_count)
    meta = PngImagePlugin.PngInfo()
    meta.add_text("sd_json", json.dumps(metadata))
    img.save(fname, pnginfo=meta)
    with open(fname + ".json", "w", encoding="utf-8") as f:
        json.dump(metadata, f, indent=2)
    return fname


def call_stable_diffusion(prompt_json: dict, user_input: str) -> list[str]:
    """
    Generate 'n_imgs' images from diffusers, save to disk with embedded prompt JSON,
    write sidecar .json, and return the file paths.
    """
    if torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"

    # Load pipeline
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.float32,
        safety_checker=None,
        local_files_only=True
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to(device)
    pipe.set_progress_bar_config(disable=True)
    pipe.enable_attention_slicing()

    pos_prompt = prompt_json["positive prompt"]
    neg_prompt = prompt_json["negative prompt"]
    width, height = prompt_json["dimensions"]
    steps = prompt_json.get("steps", 30)
    guidance_scale = prompt_json.get("cfg_scale", 6.5)
    batch_count = int(prompt_json.get("batch_count", 4))  # Standard: 4

    images = pipe(
        prompt=[pos_prompt] * batch_count,
        negative_prompt=[neg_prompt] * batch_count,
        width=width,
        height=height,
        num_inference_steps=steps,
        guidance_scale=guidance_scale
    ).images

    img_base = sanitize_filename(user_input)
    out_files = []
    import time
    for idx, img in enumerate(images, start=1):
        filename = get_next_available_filename(f"{img_base}-{idx}")

        # Timestamp und batch_count speichern
        prompt_json_with_time = dict(prompt_json)
        prompt_json_with_time["created"] = int(time.time())
        prompt_json_with_time["batch_count"] = batch_count

        # Metadata und Sidecar wie gehabt
        meta = PngImagePlugin.PngInfo()
        meta.add_text("sd_json", json.dumps(prompt_json_with_time))
        img.save(filename, pnginfo=meta)
        sidecar_path = filename + ".json"
        with open(sidecar_path, "w", encoding="utf-8") as f:
            json.dump(prompt_json_with_time, f, indent=2)
        out_files.append(filename)
    return out_files


# ------------- JSON “Fuzzy” Validator -------------

def validate_and_fix_json(raw: str) -> tuple[dict, str] | tuple[None, None]:
    """
    1) Try a strict json.loads(raw).
    2) If that fails, attempt to fix:
       - Remove trailing commas (e.g. {"a":1,} -> {"a":1})
       - Replace single quotes with double quotes (naïve)
    3) If the “fixed” string parses, return (parsed_dict, fixed_string).
    Otherwise, return (None, None).
    """
    try:
        parsed = json.loads(raw)
        return parsed, raw
    except Exception:
        # Attempt quick fixes:
        s = raw
        # 1) Remove any trailing commas before } or ]
        s = re.sub(r",\s*([}\]])", r"\1", s)
        # 2) Replace single quotes around keys/values with double quotes (naïve)
        s = re.sub(r"(?P<pre>[:\s])'(?P<val>[^']*)'", r'\1"\g<val>"', s)
        try:
            parsed = json.loads(s)
            return parsed, s
        except Exception:
            return None, None

# ------------- PyWebview API -------------

class Api:
    def __init__(self):
        self.window = None
        self.last_object = ""
        self.last_prompt_json = None
        self._abort_sd = False
        self._placeholders_ready_events = {}

        # Keep track of file modification times
        self._mtimes: dict[str, float] = {}
        # Start folder watcher thread
        watcher = threading.Thread(target=self._monitor_folder, daemon=True)
        watcher.start()

    def set_window(self, win):
        self.window = win

    def generate_prompt(self, object_name: str):
        """
        Runs only the LLM -> JSON steps (no Stable Diffusion).
        """
        self.last_object = object_name
        self.last_prompt_json = None
        threading.Thread(target=self._run_llm_to_json, args=(object_name,), daemon=True).start()
        return {"status": "started_prompt"}

    def _run_llm_to_json(self, object_name: str):
        # 1) LLM streaming
        llm_text = ""
        llm_prompt = META_PROMPT.replace("{{OBJECT}}", object_name)
        for chunk in stream_ollama(llm_prompt):
            llm_text = chunk
            self._js("window.on_llm_output", llm_text)
        self._js("window.on_llm_done")

        # 2) JSON streaming
        json_text = ""
        json_prompt = SECOND_PROMPT_TEMPLATE.replace("{PROMPT}", llm_text)
        for chunk in get_json_from_llm(json_prompt):
            json_text = chunk
            self._js("window.on_json_output", json_text)
        self._js("window.on_json_done")

        # 3) Parse JSON
        try:
            prompt_dict = json.loads(json_text)
        except Exception as e:
            self._js("window.on_error", f"JSON parse error: {e}")
            return

        # 4) Augment with SD settings
        prompt_dict["checkpoint"] = "RealismPlus/dreamshaper_8.safetensors"
        prompt_dict["vae"] = "Automatic"
        prompt_dict["sampler"] = "DPM++ 2M Karras"
        prompt_dict["steps"] = 30
        prompt_dict["cfg_scale"] = 6.5
        prompt_dict["batch_count"] = 4

        # 5) Send full JSON
        full_json_str = json.dumps(prompt_dict, indent=2)
        self._js("window.on_json_output", full_json_str)

        # 6) Store for images
        self.last_prompt_json = prompt_dict

    def generate_images(self, json_str: str):
        """
        Called from JS when “Generate Images” is clicked.
        1) Validate/fix JSON (fuzzy).
        2) Ensure “positive prompt” is nonempty.
        3) Auto fill missing keys (dimensions, checkpoint, etc.)
        4) Kick off Stable Diffusion.
        """
        parsed, fixed = validate_and_fix_json(json_str)
        if parsed is None:
            self._js("window.on_error", "Cannot parse JSON (even after auto-fix).")
            return {"status": "invalid_json"}

        pos = parsed.get("positive prompt", "").strip()
        if not pos:
            self._js("window.on_error", "“positive prompt” must not be empty.")
            return {"status": "empty_positive"}

        # Auto-fill missing/malformed keys
        dims = parsed.get("dimensions")
        if not (
            isinstance(dims, list)
            and len(dims) == 2
            and all(isinstance(x, int) for x in dims)
        ):
            parsed["dimensions"] = [768, 768]
        parsed.setdefault("checkpoint", "RealismPlus/dreamshaper_8.safetensors")
        parsed.setdefault("vae", "Automatic")
        parsed.setdefault("sampler", "DPM++ 2M Karras")
        parsed.setdefault("steps", 30)
        parsed.setdefault("cfg_scale", 6.5)
        parsed.setdefault("batch_count", 4)
        
        # If we “fixed” it, push corrected JSON back
        full_json_str = json.dumps(parsed, indent=2)
        if fixed is None or full_json_str != json_str:
            self._js("window.on_json_output", full_json_str)

        # Store and infer last_object if not yet set
        self.last_prompt_json = parsed
        if not self.last_object:
            # Derive a base name from the positive prompt (before first comma)
            cand = pos.split(",")[0]
            base = re.sub(r"\s+", "_", cand.strip().lower())
            base = re.sub(r"[^a-z0-9_]+", "", base)
            self.last_object = base or "image"

        # Launch SD thread – vorher Abort-Flag zurücksetzen
        self._abort_sd = False
        threading.Thread(target=self._run_sd_only, daemon=True).start()
        return {"status": "started_images"}

    def _get_image_generation_backend(self) -> str:
        return normalize_image_backend(load_app_settings().get("image_generation_backend"))

    def _get_replicate_image_model(self) -> str:
        return normalize_replicate_image_model(load_app_settings().get("replicate_image_model"))

    def _run_local_diffusion_images(self, batch_count: int) -> list[Image.Image]:
        device = "mps" if torch.backends.mps.is_available() else "cpu"

        pipe = StableDiffusionPipeline.from_pretrained(
            MODEL_PATH, torch_dtype=torch.float32,
            safety_checker=None, local_files_only=True
        )
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        pipe = pipe.to(device)
        pipe.set_progress_bar_config(disable=True)
        pipe.enable_attention_slicing()

        pos = self.last_prompt_json["positive prompt"]
        neg = self.last_prompt_json["negative prompt"]
        w, h = self.last_prompt_json["dimensions"]
        steps = self.last_prompt_json.get("steps", 30)
        cfg = self.last_prompt_json.get("cfg_scale", 6.5)

        def _check_abort(*args, **kwargs):
            if self._abort_sd:
                raise RuntimeError("User aborted")
            return {}

        outputs = pipe(
            prompt=[pos] * batch_count,
            negative_prompt=[neg] * batch_count,
            width=w, height=h,
            num_inference_steps=steps,
            guidance_scale=cfg,
            callback_on_step_end=_check_abort,
        )
        return outputs.images

    def _build_replicate_image_input(self, model_id: str, num_outputs: int) -> dict:
        model = REPLICATE_IMAGE_MODELS[model_id]
        w, h = self.last_prompt_json["dimensions"]
        ratio = nearest_aspect_ratio(w, h, model["aspect_ratios"])
        prompt = compose_replicate_prompt(self.last_prompt_json)

        if model_id == "black-forest-labs/flux-schnell":
            steps = int(self.last_prompt_json.get("steps", 4))
            return {
                "prompt": prompt,
                "aspect_ratio": ratio,
                "num_outputs": max(1, min(int(num_outputs), 4)),
                "num_inference_steps": max(1, min(steps, 4)),
                "output_format": "png",
                "go_fast": True,
                "megapixels": "1",
            }

        if model_id == "google/imagen-4-fast":
            return {
                "prompt": prompt,
                "aspect_ratio": ratio,
                "output_format": "png",
                "safety_filter_level": "block_only_high",
            }

        if model_id == "bytedance/seedream-4":
            return {
                "prompt": prompt,
                "aspect_ratio": ratio,
                "size": "2K",
                "max_images": 1,
                "enhance_prompt": False,
                "sequential_image_generation": "disabled",
            }

        raise RuntimeError(f"Unsupported Replicate image model: {model_id}")

    def _run_replicate_prediction(self, model_id: str, input_payload: dict):
        token = self._get_replicate_api_token()
        if not token:
            raise RuntimeError("Replicate API token missing. Open settings and add it.")

        client = replicate.Client(api_token=token)
        client.poll_interval = REPLICATE_IMAGE_POLL_INTERVAL
        prediction = client.models.predictions.create(
            model=model_id,
            input=input_payload,
            wait=False,
        )
        print(f"[Replicate Image] Prediction started: {prediction.id} ({prediction.status})")

        deadline = time.monotonic() + REPLICATE_IMAGE_TIMEOUT
        last_status = prediction.status
        while prediction.status not in {"succeeded", "failed", "canceled"}:
            if self._abort_sd:
                try:
                    prediction.cancel()
                except Exception:
                    pass
                raise RuntimeError("User aborted")

            if time.monotonic() >= deadline:
                raise TimeoutError(
                    f"Timed out waiting for Replicate image prediction {prediction.id} "
                    f"after {REPLICATE_IMAGE_TIMEOUT // 60} minutes."
                )

            time.sleep(client.poll_interval)
            prediction.reload()
            if prediction.status != last_status:
                print(f"[Replicate Image] {prediction.id}: {prediction.status}")
                last_status = prediction.status

        if prediction.status != "succeeded":
            detail = prediction.error or prediction.logs or f"status={prediction.status}"
            raise RuntimeError(f"Replicate image prediction {prediction.id} failed: {detail}")

        return prediction.output

    def _run_replicate_images(self, batch_count: int) -> list[Image.Image]:
        model_id = self._get_replicate_image_model()
        images: list[Image.Image] = []

        if model_id == "black-forest-labs/flux-schnell":
            remaining = batch_count
            while remaining > 0:
                if self._abort_sd:
                    raise RuntimeError("User aborted")
                chunk = min(remaining, 4)
                output = self._run_replicate_prediction(
                    model_id,
                    self._build_replicate_image_input(model_id, chunk),
                )
                urls = extract_replicate_output_urls(output)
                if len(urls) < chunk:
                    raise RuntimeError(f"Replicate returned {len(urls)} image(s), expected {chunk}.")
                images.extend(load_image_from_url(url) for url in urls[:chunk])
                remaining -= chunk
            return images

        for _ in range(batch_count):
            if self._abort_sd:
                raise RuntimeError("User aborted")
            output = self._run_replicate_prediction(
                model_id,
                self._build_replicate_image_input(model_id, 1),
            )
            urls = extract_replicate_output_urls(output)
            if not urls:
                raise RuntimeError("Replicate did not return an image URL.")
            images.append(load_image_from_url(urls[0]))

        return images

    def _run_sd_only(self):
        if not self.last_prompt_json:
            self._js("window.on_error", "No valid JSON prompt available.")
            return

        batch_id    = random.randint(1_000_000, 9_999_999)
        self._last_batch_id = batch_id
        batch_count = int(self.last_prompt_json.get("batch_count", 4))

        # 1) Bail out if user already hit “Stop”
        if self._abort_sd:
            return

        # 2) Show placeholders in the UI
        self._js("window.show_placeholders", batch_count, batch_id)
        event = threading.Event()
        self._placeholders_ready_events[batch_id] = event

        # 3) Wait until JS signals placeholders are in place
        while not event.is_set():
            time.sleep(0.02)
        if self._abort_sd:
            self._js("window.remove_placeholders", batch_id)
            return

        # 4) Set up and run the selected image generator.
        try:
            if self._get_image_generation_backend() == IMAGE_BACKEND_REPLICATE:
                images = self._run_replicate_images(batch_count)
            else:
                images = self._run_local_diffusion_images(batch_count)

        except RuntimeError as e:
            # user‐initiated abort
            if str(e) == "User aborted":
                self._js("window.on_generation_aborted")
                self._js("window.remove_placeholders", batch_id)
                return
            # or some other error
            self._js("window.remove_placeholders", batch_id)
            self._js("window.on_error", f"Image generation error: {e}")
            return
        except Exception as e:
            self._js("window.remove_placeholders", batch_id)
            self._js("window.on_error", f"Image generation error: {e}")
            return

        # 5) Replace placeholders one by one
        for idx, img in enumerate(images):
            if self._abort_sd:
                self._js("window.remove_placeholders", batch_id)
                return

            fname = save_generated_image(
                img,
                self.last_prompt_json,
                self.last_object,
                idx + 1,
                batch_count,
            )

            rel = os.path.relpath(fname, WEB_DIR)
            self._js("window.replace_placeholder", idx, rel, batch_id)

        # 6) Tear down
        self._js("window.remove_placeholders", batch_id)
        self._js("window.on_image_gen_done")
        del self._placeholders_ready_events[batch_id]
        
    def stop_generation(self):
        self._abort_sd = True
        # Do NOT remove placeholders here—wait until the abort is processed
        return {"status": "aborting"}

    def placeholders_ready(self, batch_id):
        if batch_id in self._placeholders_ready_events:
            self._placeholders_ready_events[batch_id].set()
        return {"status": "acknowledged"}

    def get_settings(self):
        settings = load_app_settings()
        settings_token = str(settings.get("replicate_api_token", "")).strip()
        env_token = os.environ.get("REPLICATE_API_TOKEN", "").strip()
        source = "settings" if settings_token else "environment" if env_token else ""
        return {
            "has_replicate_api_token": bool(settings_token or env_token),
            "replicate_api_token_source": source,
            "replicate_api_token": "",
            "image_generation_backend": normalize_image_backend(settings.get("image_generation_backend")),
            "replicate_image_model": normalize_replicate_image_model(settings.get("replicate_image_model")),
            "replicate_image_models": [
                {
                    "id": model_id,
                    "label": model["label"],
                    "price": model["price"],
                }
                for model_id, model in REPLICATE_IMAGE_MODELS.items()
            ],
        }

    def save_settings(self, settings: dict):
        current = load_app_settings()
        token = str(settings.get("replicate_api_token", "")).strip()
        if token:
            current["replicate_api_token"] = token

        current["image_generation_backend"] = normalize_image_backend(
            settings.get("image_generation_backend")
        )
        current["replicate_image_model"] = normalize_replicate_image_model(
            settings.get("replicate_image_model")
        )
        save_app_settings(current)
        return {
            "status": "saved",
            "has_replicate_api_token": bool(self._get_replicate_api_token()),
            "image_generation_backend": current["image_generation_backend"],
            "replicate_image_model": current["replicate_image_model"],
        }

    def clear_replicate_api_token(self):
        current = load_app_settings()
        current.pop("replicate_api_token", None)
        save_app_settings(current)
        env_token = os.environ.get("REPLICATE_API_TOKEN", "").strip()
        return {
            "status": "cleared",
            "has_replicate_api_token": bool(env_token),
            "replicate_api_token_source": "environment" if env_token else "",
            "image_generation_backend": normalize_image_backend(current.get("image_generation_backend")),
            "replicate_image_model": normalize_replicate_image_model(current.get("replicate_image_model")),
        }

    def _get_replicate_api_token(self):
        settings_token = str(load_app_settings().get("replicate_api_token", "")).strip()
        return settings_token or os.environ.get("REPLICATE_API_TOKEN", "").strip()

    def _subprocess_env(self):
        env = os.environ.copy()
        token = self._get_replicate_api_token()
        if token:
            env["REPLICATE_API_TOKEN"] = token
        return env


    def _open_blender_with_scene(self, glb_path, hdri_path):
        # Kopiere ggf. vorbereitete Blender-Template-Datei, z.B. "scene_template.blend"
        blender_path = "/Applications/Blender.app/Contents/MacOS/Blender"
        # Template: In Blender musst du ein Skript haben, das das .glb und das hdri_path lädt
        # oder als Startscript mitgibst!
        cmd = [
            "/Applications/Blender.app/Contents/MacOS/Blender",
            "--python", resource_path("scene_setup.py"),
            "--", glb_path, hdri_path or ""
        ]
        try:
            subprocess.Popen(cmd)
            return {"status": "opened_blender"}
        except Exception as e:
            self._js("window.on_error", f"Failed to open in Blender: {e}")
            return {"status": "blender_error"}


    def edit_external(self, filepath: str):
        # filepath kommt von 2D als "output/foo.png" oder von 3D als "output/foo.png"
        ext = os.path.splitext(filepath)[1].lower()
        full_path = os.path.join(WEB_DIR, filepath)
        # Ist es ein Bild oder das PNG eines 3D-Modells?
        if ext == ".png":
            # Prüfe, ob ein GLB daneben liegt
            fbase = os.path.splitext(full_path)[0]
            glb_path = fbase + ".glb"
            # Prüfe, ob auch eine HDRI da ist (Sidecar JSON auslesen)
            hdri_path = None
            sidecar = full_path + ".json"
            if os.path.isfile(sidecar):
                with open(sidecar, "r", encoding="utf-8") as f:
                    meta = json.load(f)
                    if "hdri_seamless" in meta:
                        hdri_path = os.path.join(os.path.dirname(full_path), meta["hdri_seamless"])
            # Wenn .glb existiert → ist 3D Model: Blender öffnen!
            if os.path.isfile(glb_path):
                return self._open_blender_with_scene(glb_path, hdri_path)
            # Sonst wie bisher (Photoshop)
            try:
                subprocess.Popen(["open", "-a", "Adobe Photoshop 2024", full_path])
            except Exception as e:
                self._js("window.on_error", f"Failed to open in Photoshop: {e}")
            return {"status": "opened_external"}


            # ------------- Generate 3D Model -------------
    def generate_3d_model(self, filepath: str):
        # filepath kommt aus JS als "output/flower-1.png"
        full_path = os.path.join(WEB_DIR, filepath)
        if not os.path.isfile(full_path):
            self._js("window.on_error", f"Image not found: {full_path}")
            return {"status": "missing_file"}

        threading.Thread(target=self._run_generate_3d, args=(full_path,), daemon=True).start()
        return {"status": "started_3d"}

    def _run_generate_3d(self, img_path: str):
        """
        Execute the user‐provided script `image_to_3d.py <img_path>`, wait for completion,
        then, if a .glb was produced, call JS callback window.on_3d_generated(glb_path).
        Otherwise, call window.on_error(…).
        """
        try:
            if not self._get_replicate_api_token():
                self._js(
                    "window.on_error",
                    "Replicate API token missing. Open settings and add it.",
                )
                return

            glb_path = self._run_generate_3d_and_return(img_path)
            self._js("window.on_3d_generated", glb_path)
        except Exception as e:
            self._js("window.on_error", f"3D conversion exception: {e}")





            # ------------- NEW: Generate 3D Model PLUS equirectangular map -------------

    def generate_3d_and_hdri(self, filepath: str):
        """
        Called when user wants both 3D model and equirectangular HDRI for an image.
        Kicks off both processes, waits for both, and (later) calls the next step.
        """
        full_path = os.path.join(WEB_DIR, filepath)
        if not os.path.isfile(full_path):
            self._js("window.on_error", f"Image not found: {full_path}")
            return {"status": "missing_file"}
        # Start thread for orchestration
        threading.Thread(target=self._run_3d_and_hdri, args=(full_path,), daemon=True).start()
        return {"status": "started_3d_and_hdri"}

    def _run_3d_and_hdri(self, img_path: str):
        # 1. Get object info (from PNG metadata or sidecar)
        obj_name = self._extract_object_from_json(img_path)
        print(f"[3D+HDRI] Object name from JSON: {obj_name}")
        if not obj_name:
            self._js("window.on_error", "Could not determine object name for HDRI prompt.")
            return

        # 2. Generate short background prompt using Mistral (Ollama) with streaming debug
        user_prompt = f"""Suggest a short, vivid background description for a 3D scene featuring {obj_name}. \
    The background should be immersive and plausible, suitable as an equirectangular HDRI for 3D rendering. \
    Do NOT mention the object itself. Describe the environment in a concise way, e.g.: "lush wheat field under a blue sky", "bustling medieval town square at dawn", "mysterious alien jungle with glowing plants", etc. Respond only with the description, no extra text. Only describe the appearance of the environment, no other attributes than what is to see. \
    Keep it short and precise, use 2 to 5 words max! No punctuation, just commas if needed. All lowercase. As short as possible!"""
        print(f"[HDRI LLM] Sending prompt to LLM:\n{user_prompt}\n")

        pano_prompt = None
        for chunk in stream_ollama(user_prompt):
            print(f"[HDRI LLM] {chunk!r}")
            pano_prompt = chunk.strip()
        if not pano_prompt:
            self._js("window.on_error", "Could not generate panorama prompt.")
            return

        # 3. Prepend required prefix and print
        pano_prompt_full = "hdri panorama view, equirectangular, 3d rendering of " + pano_prompt
        prompt_with_colons = f":{pano_prompt_full}:"
        print(f"[HDRI PROMPT] Final panorama prompt:\n{prompt_with_colons}\n")

        # 4. Prepare output path in OUTPUT_DIR
        out_base = os.path.splitext(os.path.basename(img_path))[0]
        hdri_filename = f"{out_base}_hdri_seamless.png"
        hdri_output_path = os.path.join(OUTPUT_DIR, hdri_filename)
        print(f"[HDRI PATH] Will save to: {hdri_output_path}")

        # 5. Run both tasks concurrently
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor() as executor:
            fut_3d = executor.submit(self._run_generate_3d_and_return, img_path)
            fut_hdri = executor.submit(self._run_equirect_map, prompt_with_colons, hdri_output_path)
            done, _ = concurrent.futures.wait([fut_3d, fut_hdri], return_when=concurrent.futures.ALL_COMPLETED)

        # 6. Check results
        err3d = fut_3d.exception() if fut_3d.done() else "3D model task did not finish"
        errhdri = fut_hdri.exception() if fut_hdri.done() else "HDRI task did not finish"

        if err3d or errhdri:
            msg = ""
            if err3d:
                msg += f"3D model error: {err3d}\n"
            if errhdri:
                msg += f"HDRI error: {errhdri}\n"
            self._js("window.on_error", msg.strip())
            return

        glb_path = fut_3d.result()
        hdri_path = fut_hdri.result()

        # 7. Write HDRI path into the original image's JSON sidecar (absolute, relative, or both)
        sidecar_path = img_path + ".json"
        try:
            if os.path.isfile(sidecar_path):
                with open(sidecar_path, "r", encoding="utf-8") as f:
                    meta = json.load(f)
            else:
                meta = {}

            meta["hdri_seamless"] = os.path.relpath(hdri_path, OUTPUT_DIR)  # or just hdri_filename
            with open(sidecar_path, "w", encoding="utf-8") as f:
                json.dump(meta, f, indent=2)
            print(f"[HDRI JSON] Wrote HDRI path to {sidecar_path}: {meta['hdri_seamless']}")
        except Exception as e:
            print(f"[HDRI JSON] ERROR writing HDRI path: {e}")

        self._on_3d_and_hdri_ready(glb_path, hdri_path)

    def _extract_object_from_json(self, img_path: str) -> str:
        """
        Helper: Tries to read the object/positive prompt from sidecar or PNG metadata.
        """
        try:
            sidecar = img_path + ".json"
            if os.path.isfile(sidecar):
                with open(sidecar, "r", encoding="utf-8") as f:
                    meta = json.load(f)
                    return meta.get("positive prompt", "").split(",")[0]
            else:
                # Try embedded PNG metadata as fallback
                img = Image.open(img_path)
                info = img.info
                if "sd_json" in info:
                    d = json.loads(info["sd_json"])
                    return d.get("positive prompt", "").split(",")[0]
        except Exception:
            pass
        return None

    def _get_panorama_prompt(self, obj_name: str) -> str:
        """
        Calls Mistral (Ollama) to generate a short background panorama prompt.
        """
        user_prompt = f"""Suggest a short, vivid background description for a 3D scene featuring {obj_name}. \
The background should be immersive and plausible, suitable as an equirectangular HDRI for 3D rendering. \
Do NOT mention the object itself. Describe the environment in a concise way, e.g.: "a lush wheat field under a blue sky", "a bustling medieval town square at dawn", "a mysterious alien jungle with glowing plants", etc. Respond only with the description, no extra text."""

        try:
            for response in stream_ollama(user_prompt):
                last = response.strip()
            return last
        except Exception:
            return None

    def _run_generate_3d_and_return(self, img_path: str):
        """
        Same as _run_generate_3d, but returns the GLB path on success (or raises).
        """
        if not self._get_replicate_api_token():
            raise RuntimeError("Replicate API token missing. Open settings and add it.")

        import image_to_3d

        result = image_to_3d.process_image(img_path, self._get_replicate_api_token())
        if result and os.path.isfile(result):
            return result

        raise RuntimeError("3D conversion did not produce .glb")

    def _run_equirect_map(self, prompt: str, output_path: str):
        """
        Calls your equi map generator script.
        Logs command and full output for debugging.
        """
        import generate_equirect

        print(f"[HDRI] Generating to: {output_path}")
        result = generate_equirect.generate_equirect(
            prompt,
            output_path,
            work_dir=APP_SUPPORT_DIR if IS_FROZEN else BASE_DIR,
        )
        if result and os.path.isfile(output_path):
            print(f"[HDRI OK] Created: {output_path}")
            return output_path
        else:
            raise RuntimeError(f"HDRI generation failed for {output_path}")

    def _on_3d_and_hdri_ready(self, glb_path: str, hdri_path: str):
        """
        Dummy function to handle successful completion of both tasks.
        For now: just print/log, later: implement what you need (e.g. combine assets, notify UI, etc).
        """
        print(f"3D model ready: {glb_path}")
        print(f"HDRI ready: {hdri_path}")
        # Example: self._js("window.on_3d_and_hdri_ready", glb_path, hdri_path)













    # ------------- END of Generate 3D Model & Equi methods -------------
    def get_initial_images(self):
        entries = []
        for fname in os.listdir(OUTPUT_DIR):
            if not fname.lower().endswith(".png"):
                continue

            fpath = os.path.join(OUTPUT_DIR, fname)
            # Prüfe, ob Sidecar existiert & valides JSON ist
            sidecar = fpath + ".json"
            if not os.path.isfile(sidecar):
                continue
            try:
                with open(sidecar, "r", encoding="utf-8") as f:
                    meta = json.load(f)
                    # Optional: Check ob wirklich ein Prompt drin ist
                    if not meta.get("positive prompt"):
                        continue
                    created = int(meta.get("created", 0))
            except Exception:
                continue

            rel = os.path.relpath(os.path.join(OUTPUT_DIR, fname), WEB_DIR)
            entries.append({
                "filepath": rel,
                "created": created
            })

        return sorted(entries, key=lambda e: e["created"], reverse=True)

    def get_3d_models(self):
        entries = []
        for fname in os.listdir(OUTPUT_DIR):
            if not fname.lower().endswith(".png"):
                continue
            png_path = os.path.join(OUTPUT_DIR, fname)
            fbase = os.path.splitext(fname)[0]
            glb_path = os.path.join(OUTPUT_DIR, fbase + ".glb")
            if os.path.isfile(glb_path):
                # Lade "created" aus Sidecar-JSON
                created = 0
                sidecar = png_path + ".json"
                if os.path.isfile(sidecar):
                    try:
                        with open(sidecar, "r", encoding="utf-8") as f:
                            meta = json.load(f)
                            created = int(meta.get("created", 0))
                    except Exception:
                        pass
                rel_img = os.path.relpath(png_path, WEB_DIR)
                rel_glb = os.path.relpath(glb_path, WEB_DIR)
                entries.append({
                    "img": rel_img,
                    "glb": rel_glb,
                    "created": created
                })
        # Sortiere nach created DESC
        models = sorted(entries, key=lambda e: e["created"], reverse=True)
        # Entferne "created" vor dem Rückgeben (falls Frontend das nicht braucht)
        for m in models:
            if "created" in m:
                del m["created"]
        return models
        
    def get_image_json(self, filepath: str):
        # filepath kommt aus JS als "output/flower-1.png"
        full_path = os.path.join(WEB_DIR, filepath)
        try:
            # Always prefer the sidecar if available (contains most up-to-date data)
            sidecar = full_path + ".json"
            if os.path.isfile(sidecar):
                with open(sidecar, "r", encoding="utf-8") as f:
                    d = json.load(f)
                    if "created" in d:
                        del d["created"]
                    return json.dumps(d, indent=2)
            # Fallback: try embedded metadata in PNG
            img = Image.open(full_path)
            info = img.info
            if "sd_json" in info:
                d = json.loads(info["sd_json"])
                if "created" in d:
                    del d["created"]
                return json.dumps(d, indent=2)
            else:
                return json.dumps({"error": "No embedded generation data found."})
        except Exception as e:
            return json.dumps({"error": f"Failed to read metadata: {e}"})

    def _monitor_folder(self):
        """
        Every 2 seconds, scan OUTPUT_DIR for PNGs. If any file is new or its mtime changed,
        send window.on_image_updated(filepath, dataUri) to front end.
        """
        while True:
            time.sleep(2)
            for fname in os.listdir(OUTPUT_DIR):
                if not fname.lower().endswith(".png"):
                    continue
                fpath = os.path.join(OUTPUT_DIR, fname)
                try:
                    mtime = os.path.getmtime(fpath)
                except FileNotFoundError:
                    continue
                previous = self._mtimes.get(fpath)
                if previous is None or mtime > previous:
                    # New or modified
                    self._mtimes[fpath] = mtime
                    try:
                        data_uri = image_to_base64(fpath)
                        rel = os.path.relpath(fpath, WEB_DIR)  # → "output/flower-2.png"
                        self._js("window.on_image_updated", rel, data_uri)
                    except Exception:
                        pass

    def _js(self, fn: str, *args):
        import json
        args_json = json.dumps(list(args))
        code = f"{fn}.apply(null, {args_json})"
        self.window.evaluate_js(code)


# ------------- Launch PyWebview -------------

if __name__ == "__main__":
    api = Api()
    window = webview.create_window(
        "SD 3D Model Image Gen",
        WEB_INDEX,
        js_api=api,
        width=900,
        height=1000,
        min_size=(650, 500),
        resizable=True
    )
    api.set_window(window)
    webview.start(debug=not IS_FROZEN, icon=APP_ICON_PATH if os.path.isfile(APP_ICON_PATH) else None)