import gradio as gr
from daggr import GradioNode, InferenceNode, FnNode, Graph

# 1. Text generation with Llama 3
llama3 = InferenceNode(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    inputs={
        "prompt": gr.Textbox(label="Your request"),
        "max_tokens": 500,
        "temperature": 0.7
    },
    outputs={
        "response": gr.Textbox(label="Generated Text")
    },
    name="Llama 3 Text Generator"
)

# 2. Image generation from text
image_gen = GradioNode(
    space_or_url="stabilityai/stable-diffusion-xl-base-1.0",
    api_name="/predict",
    inputs={
        "prompt": gr.Textbox(label="Image prompt"),
        "negative_prompt": "blurry, low quality",
        "guidance_scale": 7.5,
        "num_inference_steps": 25
    },
    outputs={
        "image": gr.Image(label="Generated Image")
    },
    name="Stable Diffusion XL"
)

# 3. Summarization function
def summarize(text: str, max_length: int = 150) -> str:
    """Summarize text to specified length"""
    if len(text) <= max_length:
        return text
    return text[:max_length].rsplit(' ', 1)[0] + "..."

summarizer = FnNode(
    fn=summarize,
    inputs={
        "text": gr.Textbox(label="Text to summarize"),
        "max_length": gr.Slider(50, 300, value=150, label="Summary length")
    },
    outputs={
        "summary": gr.Textbox(label="Summary")
    },
    name="Text Summarizer"
)

# 4. Combine nodes into a workflow
workflow = Graph(
    name="AI Content Workflow",
    nodes=[llama3, image_gen, summarizer],
    connections=[
        (llama3.outputs["response"], image_gen.inputs["prompt"]),
        (llama3.outputs["response"], summarizer.inputs["text"])
    ]
)

# Launch the application
workflow.launch()