app.py

import base64
import json
import os
import random
import re
import shutil
import sys
import tempfile
import uuid
import requests
from datetime import datetime
from io import BytesIO
from pathlib import Path

import gradio as gr
from PIL import Image
from dotenv import load_dotenv
from graphviz import Digraph
from huggingface_hub import InferenceClient
from together import Together

# ────────────────────────────────
# ENV / API
# ────────────────────────────────
load_dotenv()
HF_TOKEN        = os.getenv("HF_TOKEN")            # <-- add your HF token to .env
TOGETHER_TOKEN  = os.getenv("TOGETHER_API_KEY", "")
together_client = Together(api_key=TOGETHER_TOKEN)
image_client    = InferenceClient(token=HF_TOKEN)  # default model set later

# Optional Graphviz path helper (Windows ONLY (RIP Gotham))
# if shutil.which("dot") is None:
#     gv_path = r"C:\Program Files\Graphviz\bin"
#     if os.path.exists(gv_path):
#         os.environ["PATH"] = gv_path + os.pathsep + os.environ["PATH"]
#     else:
#         sys.exit("Graphviz not found. Please install Graphviz or remove the check.")

# ────────────────────────────────
# LLM templates
# ────────────────────────────────
LLAMA_JSON_PROMPT = """
Extract every character and any explicit relationship between them.
Return pure JSON ONLY in this schema:
{
  "characters": ["Alice", "Bob"],
  "relations": [
     {"from":"Alice","to":"Bob","type":"friend"}
  ]
}
TEXT:
\"\"\"%s\"\"\"
"""

IMAGE_PROMPT_TEMPLATE = """
Based on the following story, write %d distinct vivid scene descriptions, one per line.
Each line should begin with a dash (-) followed by a detailed image-worthy scene.
Include setting, mood, characters, and visual cues.

Return ONLY the list of scenes, each on its own line.

Story:
\"\"\"%s\"\"\"
"""

# ────────────────────────────────
# Entity extraction
# ────────────────────────────────
def extract_entities(text: str):
    try:
        prompt = LLAMA_JSON_PROMPT % text
        resp = together_client.chat.completions.create(
            model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1024,
        )
        raw = resp.choices[0].message.content.strip()
        m = re.search(r"\{[\s\S]*\}", raw)
        if not m:
            return None, f"⚠️ No JSON block found.\n\n{raw}"
        data = json.loads(m.group(0))
        return data, None
    except Exception as e:
        return None, f"⚠️ extractor error: {e}"

# ────────────────────────────────
# Build visual prompt
# ────────────────────────────────
def generate_image_prompts(story_text: str, count=1):
    try:
        prompt_msg = IMAGE_PROMPT_TEMPLATE % (count, story_text)
        resp = together_client.chat.completions.create(
            model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
            messages=[{"role": "user", "content": prompt_msg}],
            max_tokens=200,
        )
        raw_output = resp.choices[0].message.content.strip()
        prompts = [line.strip("-• ").strip() for line in raw_output.split("\n") if line.strip()]
        return prompts[:count]  # just in case LLM gives more than needed
    except Exception as e:
        print("⚠️ LLM scene prompt generation failed:", e)
        return []



def generate_images_with_together(story, style, quality, count=1):
    base_prompt = generate_image_prompts(story)
    images = []

    for i in range(count):
        full_prompt = f"{style} style, cinematic lighting, quality {quality}, {base_prompt} [Scene {i + 1}]"
        seed = random.randint(1, 10_000_000)

        try:
            resp = together_client.images.generate(
                model="black-forest-labs/FLUX.1-schnell-Free",
                prompt=full_prompt,
                seed=seed,
                width=768,
                height=512,
                steps=4
            )
        except Exception as e:
            print("🔥 Together image API error:", e)
            break

        img = None
        if resp.data:
            choice = resp.data[0]

            if getattr(choice, "url", None):
                try:
                    img_bytes = requests.get(choice.url, timeout=30).content
                    img = Image.open(BytesIO(img_bytes))
                except Exception as e:
                    print("⚠️ URL fetch failed:", e)

            elif getattr(choice, "b64_json", None):
                try:
                    img_bytes = base64.b64decode(choice.b64_json)
                    img = Image.open(BytesIO(img_bytes))
                except Exception as e:
                    print("⚠️ base64 decode failed:", e)

        if img is not None:
            images.append(img)
        else:
            print(f"⚠️ No image for scene {i+1}")

    return images

# ────────────────────────────────
# Graph → PNG (Graphviz)
# ────────────────────────────────
def build_graph_png(data: dict) -> str:
    dot = Digraph(format="png")
    dot.attr(rankdir="LR", bgcolor="white", fontsize="11")
    for c in data["characters"]:
        dot.node(c, shape="ellipse", style="filled", fillcolor="#8ecae6")
    for r in data["relations"]:
        dot.edge(r["from"], r["to"], label=r["type"], fontsize="10")

    tmpdir = Path(tempfile.mkdtemp())
    path = tmpdir / f"graph_{uuid.uuid4().hex}.png"
    dot.render(path.stem, directory=tmpdir, cleanup=True)
    return str(path)

# ────────────────────────────────
# Core generation
# ────────────────────────────────
def generate_assets(prompt, style, quality, num_images, state):
    data, err = extract_entities(prompt)
    if not data:
        return [], None, err or "No data.", state

    graph_path = build_graph_png(data)

    images = []
    if num_images > 0:
        try:
            images = generate_images_with_together(prompt, style, quality, int(num_images))
        except Exception as e:
            status = f"⚠️ Image generation failed: {e}"
            return [], graph_path, status, data

    status = "✅ All assets generated." if images else "✅ Graph generated (no images)."
    return images, graph_path, status, data

# Helper to rebuild graph after manual edits
def _regen_graph(state): return gr.update(value=build_graph_png(state))

# ────────────────────────────────
# Manual tweak callbacks
# ────────────────────────────────
def add_character(name, state):
    if not name:
        return gr.update(), "Enter a character name.", state
    if name in state["characters"]:
        return gr.update(), f"{name} already exists.", state
    state["characters"].append(name)
    return _regen_graph(state), "✅ Character added.", state

def add_relation(frm, to, typ, state):
    if frm not in state["characters"] or to not in state["characters"]:
        return gr.update(), "Both characters must exist first.", state
    state["relations"].append({"from": frm, "to": to, "type": typ or "relation"})
    return _regen_graph(state), "✅ Relation added.", state

def delete_character(name, state):
    if name not in state["characters"]:
        return gr.update(), "Character not found.", state
    state["characters"].remove(name)
    state["relations"] = [r for r in state["relations"] if r["from"] != name and r["to"] != name]
    return _regen_graph(state), f"🚮 {name} deleted.", state

# Save / Load
def save_json(state):
    fp = Path(tempfile.gettempdir()) / f"story_{datetime.now().isoformat()}.json"
    fp.write_text(json.dumps(state, indent=2))
    return str(fp)

def load_json(file_obj, state):
    if not file_obj or not Path(file_obj).exists():
        return gr.update(), "No file uploaded.", state
    try:
        data = json.loads(Path(file_obj).read_text())
        assert "characters" in data and "relations" in data
        return _regen_graph(data), "✅ File loaded.", data
    except Exception as e:
        return gr.update(), f"Load error: {e}", state

# ────────────────────────────────
# UI (same tabs you designed)
# ────────────────────────────────
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan")) as demo:
    gr.Markdown("## ✨ EpicFrame – Narrative Workbench")

    state = gr.State({"characters": [], "relations": []})

    # Input tab
    with gr.Tab("Input"):
        text_input     = gr.Textbox(label="Story prompt", lines=6)
        style_dropdown = gr.Dropdown(["Realistic", "Anime", "Sketch"], value="Realistic", label="Style")
        quality_slider = gr.Slider(1, 10, value=7, step=1, label="Image Quality")
        num_images_sl  = gr.Slider(0, 4, value=0, step=1, label="Images to generate (0 = skip)")
        generate_btn   = gr.Button("▶️ Generate Assets")
        status_box     = gr.Textbox(label="Status", lines=2)

    # Images tab
    with gr.Tab("Images"):
        gallery = gr.Gallery(label="🖼️ Images", columns=4)

    # Graph/Edit tab
    with gr.Tab("Graph / Edit"):
        graph_img = gr.Image(label="📌 Character Map", interactive=False, height=500)
        with gr.Row():
            add_char_name = gr.Textbox(label="Add Character → Name")
            add_char_btn  = gr.Button("Add")
        with gr.Row():
            rel_from  = gr.Textbox(label="Relation From")
            rel_to    = gr.Textbox(label="To")
            rel_type  = gr.Textbox(label="Type")
            add_rel_btn = gr.Button("Add Relation")
        with gr.Row():
            del_char_name = gr.Textbox(label="Delete Character → Name")
            del_char_btn  = gr.Button("Delete")
        tweak_msg = gr.Textbox(label="➰ Status", max_lines=2)

    # Save/Load tab
    with gr.Tab("Save / Load"):
        save_btn   = gr.Button("💾 Download JSON")
        load_file  = gr.File(label="Load JSON")
        load_btn   = gr.Button("⤵️ Load into workspace")
        save_msg   = gr.Textbox(label="Status", max_lines=2)

    # callbacks
    generate_btn.click(
        generate_assets,
        inputs=[text_input, style_dropdown, quality_slider, num_images_sl, state],
        outputs=[gallery, graph_img, status_box, state]
    )

    add_char_btn.click(add_character,
                       inputs=[add_char_name, state],
                       outputs=[graph_img, tweak_msg, state])
    add_rel_btn.click(add_relation,
                      inputs=[rel_from, rel_to, rel_type, state],
                      outputs=[graph_img, tweak_msg, state])
    del_char_btn.click(delete_character,
                       inputs=[del_char_name, state],
                       outputs=[graph_img, tweak_msg, state])

    save_btn.click(save_json, inputs=state, outputs=save_btn, api_name="download") \
            .then(lambda p: "✅ JSON ready.", outputs=save_msg)
    load_btn.click(load_json, inputs=[load_file, state],
                   outputs=[graph_img, save_msg, state])

demo.launch()