from __future__ import annotations

import base64
import io
import json
import os
import re
import tempfile
import time
import zipfile
from pathlib import Path
from typing import Any

import gradio as gr
from gradio_client import Client, handle_file
from huggingface_hub import InferenceClient
from PIL import Image, PngImagePlugin


APP_TITLE = "Expression Pack Studio"
BACKEND_SPACE = os.getenv("BACKEND_SPACE", "jblast94/Qwen-Image-Edit-NSFW")
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
TEXT_TO_IMAGE_MODEL = os.getenv("TEXT_TO_IMAGE_MODEL", "black-forest-labs/FLUX.1-schnell")
DEFAULT_NEGATIVE = "blurry, low quality, deformed, extra limbs, watermark, text, logo, duplicate face"


EXPRESSION_PRESETS: dict[str, dict[str, str]] = {
    "Starter 8 - fast test": {
        "neutral": "Change only the expression to a relaxed neutral face. Preserve identity, hair, outfit, pose, lighting, and background.",
        "joy": "Change only the expression to a genuine happy smile with bright eyes. Preserve identity, hair, outfit, pose, lighting, and background.",
        "sadness": "Change only the expression to soft sadness with slightly downturned eyes and mouth. Preserve everything else.",
        "anger": "Change only the expression to anger with furrowed brows and an intense stare. Preserve everything else.",
        "surprise": "Change only the expression to wide-eyed surprise with a slightly open mouth. Preserve everything else.",
        "fear": "Change only the expression to fear with worried wide eyes and a tense mouth. Preserve everything else.",
        "confused": "Change only the expression to confusion with slightly furrowed brows and uncertain eyes. Preserve everything else.",
        "smug": "Change only the expression to a smug confident smirk. Preserve everything else.",
    },
    "Standard 20": {
        "neutral": "Change only the expression to a relaxed neutral face. Preserve identity, hair, outfit, pose, lighting, and background.",
        "joy": "Change only the expression to a genuine happy smile with bright eyes. Preserve everything else.",
        "happiness": "Change only the expression to warm content happiness with a relaxed smile. Preserve everything else.",
        "sadness": "Change only the expression to soft sadness with slightly downturned eyes and mouth. Preserve everything else.",
        "sorrow": "Change only the expression to sorrow with wet eyes and a slight frown. Preserve everything else.",
        "anger": "Change only the expression to anger with furrowed brows and an intense stare. Preserve everything else.",
        "rage": "Change only the expression to rage with fierce eyes and a clenched jaw. Preserve everything else.",
        "surprise": "Change only the expression to wide-eyed surprise with a slightly open mouth. Preserve everything else.",
        "shock": "Change only the expression to shock with raised eyebrows and an open mouth. Preserve everything else.",
        "fear": "Change only the expression to fear with worried wide eyes and a tense mouth. Preserve everything else.",
        "disgust": "Change only the expression to mild disgust with a wrinkled nose and downturned mouth. Preserve everything else.",
        "calm": "Change only the expression to calm peaceful softness with relaxed eyes. Preserve everything else.",
        "confused": "Change only the expression to confusion with slightly furrowed brows and uncertain eyes. Preserve everything else.",
        "curious": "Change only the expression to curiosity with raised brows and attentive eyes. Preserve everything else.",
        "thinking": "Change only the expression to thoughtful contemplation with eyes slightly upward. Preserve everything else.",
        "serious": "Change only the expression to serious focused intensity. Preserve everything else.",
        "smug": "Change only the expression to a smug confident smirk. Preserve everything else.",
        "embarrassed": "Change only the expression to embarrassment with a slight blush and averted eyes. Preserve everything else.",
        "shy": "Change only the expression to shy softness with a gentle smile and lowered gaze. Preserve everything else.",
        "tired": "Change only the expression to tired weariness with half-closed eyes. Preserve everything else.",
    },
    "Full 28": {
        "neutral": "Change only the expression to a relaxed neutral face. Preserve identity, hair, outfit, pose, lighting, and background.",
        "joy": "Change only the expression to a genuine happy smile with bright eyes. Preserve everything else.",
        "happiness": "Change only the expression to warm content happiness with a relaxed smile. Preserve everything else.",
        "sadness": "Change only the expression to soft sadness with slightly downturned eyes and mouth. Preserve everything else.",
        "sorrow": "Change only the expression to sorrow with wet eyes and a slight frown. Preserve everything else.",
        "anger": "Change only the expression to anger with furrowed brows and an intense stare. Preserve everything else.",
        "rage": "Change only the expression to rage with fierce eyes and a clenched jaw. Preserve everything else.",
        "surprise": "Change only the expression to wide-eyed surprise with a slightly open mouth. Preserve everything else.",
        "shock": "Change only the expression to shock with raised eyebrows and an open mouth. Preserve everything else.",
        "fear": "Change only the expression to fear with worried wide eyes and a tense mouth. Preserve everything else.",
        "terror": "Change only the expression to terror with wide eyes and a strained tense face. Preserve everything else.",
        "disgust": "Change only the expression to mild disgust with a wrinkled nose and downturned mouth. Preserve everything else.",
        "contempt": "Change only the expression to contempt with a slight sneer and narrowed eyes. Preserve everything else.",
        "calm": "Change only the expression to calm peaceful softness with relaxed eyes. Preserve everything else.",
        "confused": "Change only the expression to confusion with slightly furrowed brows and uncertain eyes. Preserve everything else.",
        "curious": "Change only the expression to curiosity with raised brows and attentive eyes. Preserve everything else.",
        "thinking": "Change only the expression to thoughtful contemplation with eyes slightly upward. Preserve everything else.",
        "serious": "Change only the expression to serious focused intensity. Preserve everything else.",
        "smug": "Change only the expression to a smug confident smirk. Preserve everything else.",
        "confident": "Change only the expression to confident direct eye contact with a subtle smile. Preserve everything else.",
        "embarrassed": "Change only the expression to embarrassment with a slight blush and averted eyes. Preserve everything else.",
        "shy": "Change only the expression to shy softness with a gentle smile and lowered gaze. Preserve everything else.",
        "love": "Change only the expression to loving affection with warm eyes and a gentle smile. Preserve everything else.",
        "pout": "Change only the expression to a cute pout with pursed lips and big eyes. Preserve everything else.",
        "wink": "Change only the expression to a playful wink with one eye closed and a small smile. Preserve everything else.",
        "pain": "Change only the expression to mild pain with a slight grimace and tense brow. Preserve everything else.",
        "tired": "Change only the expression to tired weariness with half-closed eyes. Preserve everything else.",
        "bored": "Change only the expression to boredom with uninterested eyes and a flat mouth. Preserve everything else.",
    },
}


def slugify(value: str) -> str:
    value = re.sub(r"[^a-zA-Z0-9 _-]+", "", value or "character").strip()
    value = re.sub(r"\s+", "_", value)
    return value or "character"


def load_image_bytes(path: str | None) -> bytes | None:
    if not path:
        return None
    try:
        with Image.open(path) as image:
            image = image.convert("RGBA")
            buffer = io.BytesIO()
            image.save(buffer, format="PNG")
            return buffer.getvalue()
    except Exception:
        return None


def extract_card_text(path: str | None, pasted_text: str) -> str:
    pieces: list[str] = []
    if pasted_text.strip():
        pieces.append(pasted_text.strip())
    if path:
        suffix = Path(path).suffix.lower()
        if suffix in {".json", ".txt"}:
            try:
                pieces.append(Path(path).read_text(encoding="utf-8", errors="ignore"))
            except Exception:
                pass
        try:
            with Image.open(path) as image:
                info = dict(image.info)
            for key in ("chara", "ccv3", "Description", "parameters", "comment"):
                raw = info.get(key)
                if not raw:
                    continue
                text = str(raw)
                if key in {"chara", "ccv3"}:
                    try:
                        text = base64.b64decode(text).decode("utf-8", errors="ignore")
                    except Exception:
                        pass
                pieces.append(text)
        except Exception:
            pass
    return "\n\n".join(pieces)


def summarize_character(card_text: str, name: str) -> str:
    if not card_text.strip():
        return f"portrait of {name or 'a character'}, expressive face, clean character art, high quality"
    try:
        data = json.loads(card_text)
        fields = []
        for key in ("name", "description", "personality", "scenario", "first_mes"):
            value = data.get(key) or data.get("data", {}).get(key)
            if value:
                fields.append(f"{key}: {value}")
        if fields:
            return "\n".join(fields)[:2500]
    except Exception:
        pass
    return card_text[:2500]


def text_to_reference_image(character_name: str, card_text: str, style: str) -> bytes:
    if not HF_TOKEN:
        raise RuntimeError("Text-only mode needs HF_TOKEN set as a Space secret or environment variable.")
    summary = summarize_character(card_text, character_name)
    prompt = (
        f"single centered character portrait of {character_name or 'the character'}, {summary}, "
        f"{style}, detailed face, expressive eyes, clean background, high quality, no text, no watermark"
    )
    client = InferenceClient(model=TEXT_TO_IMAGE_MODEL, token=HF_TOKEN)
    image = client.text_to_image(prompt=prompt, negative_prompt=DEFAULT_NEGATIVE)
    buffer = io.BytesIO()
    image.convert("RGBA").save(buffer, format="PNG")
    return buffer.getvalue()


def write_temp_png(image_bytes: bytes) -> str:
    handle = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    handle.write(image_bytes)
    handle.close()
    return handle.name


def result_to_bytes(result: Any) -> bytes:
    candidate = result
    if isinstance(result, (list, tuple)):
        # /infer returns: result_gallery, seed, history_gallery, state.
        candidate = result[0]
    if isinstance(candidate, list):
        if not candidate:
            raise RuntimeError("Backend returned an empty result gallery.")
        candidate = candidate[0]
    if isinstance(candidate, (list, tuple)):
        if not candidate:
            raise RuntimeError("Backend returned an empty nested result gallery item.")
        candidate = candidate[0]
    if isinstance(candidate, dict):
        candidate = candidate.get("path") or candidate.get("name") or candidate.get("url")
    if isinstance(candidate, str) and candidate.startswith("http"):
        import requests

        response = requests.get(candidate, timeout=120)
        response.raise_for_status()
        return response.content
    if isinstance(candidate, str):
        return Path(candidate).read_bytes()
    raise RuntimeError(f"Unsupported image result from backend: {type(result).__name__}")


def generate_expression(reference_path: str, prompt: str, steps: int, guidance: float) -> bytes:
    client = Client(BACKEND_SPACE, hf_token=HF_TOKEN or None)
    uploaded_image = handle_file(reference_path)
    result = client.predict(
        None,                    # start state
        [uploaded_image],        # Input Images gallery
        prompt,                  # Prompt
        DEFAULT_NEGATIVE,        # Negative prompt
        0,                       # Seed
        True,                    # Randomize seed
        float(guidance),         # True guidance scale
        int(steps),              # Number of inference steps
        768,                     # Height
        768,                     # Width
        False,                   # Rewrite prompt
        0,                       # ZeroGPU max duration
        api_name="/infer",
    )
    return result_to_bytes(result)


def ensure_png(image_bytes: bytes) -> bytes:
    with Image.open(io.BytesIO(image_bytes)) as image:
        image = image.convert("RGBA")
        buffer = io.BytesIO()
        image.save(buffer, format="PNG", optimize=True)
        return buffer.getvalue()


def build_zip(character_name: str, reference: bytes, images: dict[str, bytes], metadata: dict[str, Any]) -> str:
    out_dir = Path(tempfile.mkdtemp(prefix="expression-pack-"))
    zip_path = out_dir / f"{slugify(character_name)}_expression_pack.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("reference.png", ensure_png(reference))
        for name, image_bytes in images.items():
            if image_bytes:
                zf.writestr(f"expressions/{name.lower()}.png", ensure_png(image_bytes))
                zf.writestr(f"{name.lower()}.png", ensure_png(image_bytes))
        zf.writestr("manifest.json", json.dumps(metadata, indent=2))
        zf.writestr("README.txt", "This ZIP contains root-level expression PNGs plus expressions/*.png and manifest.json.\n")
    return str(zip_path)


def generate_pack(
    reference_image: str | None,
    character_card_file: str | None,
    character_text: str,
    character_name: str,
    preset_name: str,
    custom_expressions: str,
    style: str,
    steps: int,
    guidance: float,
    progress=gr.Progress(track_tqdm=False),
):
    start = time.time()
    card_text = extract_card_text(character_card_file, character_text)
    reference = load_image_bytes(reference_image) or load_image_bytes(character_card_file)
    if reference is None:
        progress(0, desc="Creating reference portrait from text...")
        reference = text_to_reference_image(character_name, card_text, style)

    expressions = dict(EXPRESSION_PRESETS[preset_name])
    for line in custom_expressions.splitlines():
        line = line.strip()
        if not line:
            continue
        if ":" in line:
            key, value = line.split(":", 1)
            expressions[slugify(key).lower()] = value.strip()
        else:
            key = slugify(line).lower()
            expressions[key] = f"Change only the expression to {line}. Preserve identity, hair, outfit, pose, lighting, and background."

    ref_path = write_temp_png(reference)
    images: dict[str, bytes] = {}
    failures: dict[str, str] = {}
    gallery: list[tuple[bytes, str]] = []
    total = len(expressions)

    for idx, (name, prompt) in enumerate(expressions.items(), start=1):
        progress((idx - 1) / total, desc=f"Generating {idx}/{total}: {name}")
        try:
            image_bytes = generate_expression(ref_path, prompt, steps, guidance)
            images[name] = image_bytes
            gallery.append((image_bytes, name))
        except Exception as exc:
            failures[name] = str(exc)

    Path(ref_path).unlink(missing_ok=True)
    metadata = {
        "character_name": character_name or "character",
        "preset": preset_name,
        "backend_space": BACKEND_SPACE,
        "text_to_image_model": TEXT_TO_IMAGE_MODEL,
        "generated": sorted(images.keys()),
        "failed": failures,
        "seconds": round(time.time() - start, 2),
    }
    zip_path = build_zip(character_name or "character", reference, images, metadata)
    status = f"Generated {len(images)}/{total} expressions."
    if failures:
        status += f" {len(failures)} failed and were skipped; ZIP still created."
    return gallery, zip_path, json.dumps(metadata, indent=2), status


CSS = """
.gradio-container { max-width: 1180px !important; }
.hero { padding: 28px; border-radius: 24px; background: linear-gradient(135deg,#1f1235,#47216f,#a7357d); color: white; box-shadow: 0 24px 80px rgba(0,0,0,.22); }
.hero h1 { font-size: 42px; margin: 0 0 8px; }
.hero p { font-size: 16px; opacity: .92; margin: 0; }
.note { padding: 12px 16px; border-radius: 14px; background: rgba(125,85,255,.12); border: 1px solid rgba(125,85,255,.25); }
"""


with gr.Blocks(title=APP_TITLE, css=CSS, theme=gr.themes.Soft(primary_hue="purple", neutral_hue="slate")) as demo:
    gr.HTML(
        """
        <div class="hero">
            <h1>Expression Pack Studio</h1>
            <p>Upload a character photo/card or paste a description. Generate a clean downloadable ZIP of expression PNGs.</p>
        </div>
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input")
            reference_image = gr.Image(label="Reference photo / portrait", type="filepath", height=330)
            character_card_file = gr.File(label="Optional character card PNG / JSON / TXT", file_types=[".png", ".json", ".txt"], type="filepath")
            character_name = gr.Textbox(label="Character name", value="character")
            character_text = gr.Textbox(label="Character card text / description", lines=7, placeholder="Paste character description, card JSON, personality, visual details, etc.")
            style = gr.Textbox(label="Style for text-only reference generation", value="anime visual novel character portrait, polished digital art")
        with gr.Column(scale=1):
            gr.Markdown("### Pack Settings")
            preset_name = gr.Dropdown(label="Expression set", choices=list(EXPRESSION_PRESETS), value="Starter 8 - fast test")
            custom_expressions = gr.Textbox(label="Optional custom expressions", lines=8, placeholder="One per line, or name: prompt\nexample: skeptical: narrowed eyes, doubtful half-smile")
            with gr.Row():
                steps = gr.Slider(1, 12, value=4, step=1, label="Steps")
                guidance = gr.Slider(0.5, 3.0, value=1.0, step=0.1, label="Guidance")
            generate_button = gr.Button("Generate Expression Pack", variant="primary", size="lg")
            gr.Markdown("<div class='note'>Start with Starter 8. If it looks good, run Standard 20 or Full 28.</div>")
    with gr.Row():
        gallery = gr.Gallery(label="Generated expressions", columns=4, height=460, object_fit="contain")
        with gr.Column(scale=1):
            zip_file = gr.File(label="Download ZIP")
            status = gr.Textbox(label="Status", interactive=False)
            manifest = gr.Code(label="Manifest", language="json")

    generate_button.click(
        fn=generate_pack,
        inputs=[reference_image, character_card_file, character_text, character_name, preset_name, custom_expressions, style, steps, guidance],
        outputs=[gallery, zip_file, manifest, status],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))