import os
import uuid
import base64
import gradio as gr
from daggr import InferenceNode, Graph, FnNode
from daggr.state import get_daggr_files_dir
from huggingface_hub import InferenceClient

image_gen = InferenceNode(
    model="Tongyi-MAI/Z-Image",
    inputs={
        "prompt": gr.Textbox(
            label="Prompt", 
            value="A futuristic cityscape with flying cars and neon holographic advertisements, cyberpunk style", 
            lines=3
        ),
    },
    outputs={
        "image": gr.Image(label="Generated Image"),
    },
)

def summarize_image(image: str) -> str:
    client = InferenceClient(model="moonshotai/Kimi-K2.5", token=os.environ.get("HF_TOKEN"))
    
    # Encode image to base64
    with open(image, "rb") as f:
        base64_image = base64.b64encode(f.read()).decode("utf-8")
    
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "Please summarize what is happening in this image in a few sentences."},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
            ],
        }
    ]
    
    response = client.chat.completions.create(messages=messages, max_tokens=200)
    return response.choices[0].message.content

image_summary = FnNode(
    fn=summarize_image,
    inputs={
        "image": image_gen.image,
    },
    outputs={
        "summary": gr.Textbox(label="Image Summary", lines=4),
    },
)

graph = Graph(
    name="Z-Image and Kimi Summarization Workflow",
    nodes=[image_gen, image_summary]
)

graph.launch()