# app.py
from daggr import GradioNode, InferenceNode, FnNode, Graph
import gradio as gr
from typing import Dict, Any
import requests
import os

# Environment variables for API keys
API_KEYS = {
    "OPENAI": os.getenv("OPENAI_API_KEY"),
    "HUGGINGFACE": os.getenv("HF_API_KEY")
}

# ========== Input Processing Node ==========
def preprocess_inputs(user_input: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
    """Clean and validate inputs with metadata enrichment"""
    return {
        "cleaned_input": user_input.strip(),
        "timestamp": metadata.get("timestamp"),
        "source": metadata.get("source", "web")
    }

input_processor = FnNode(
    fn=preprocess_inputs,
    inputs={
        "user_input": gr.Textbox(label="User Input"),
        "metadata": gr.JSON(label="Metadata")
    },
    outputs={
        "processed_data": gr.JSON(label="Processed Input")
    }
)

# ========== LLM Processing Node ==========
llm_processor = InferenceNode(
    model="meta-llama/Llama-3-70B-Instruct",
    inputs={
        "prompt": gr.Textbox(label="LLM Prompt"),
        "temperature": gr.Slider(0, 1, value=0.7)
    },
    outputs={
        "response": gr.Textbox(label="LLM Response")
    },
    api_key=API_KEYS["HUGGINGFACE"]
)

# ========== Image Generation Node ==========
image_generator = GradioNode(
    space_or_url="stabilityai/stable-diffusion-xl-base-1.0",
    api_name="/generate",
    inputs={
        "prompt": gr.Textbox(label="Image Prompt"),
        "negative_prompt": gr.Textbox(label="Negative Prompt"),
        "steps": gr.Slider(10, 50, value=30)
    },
    outputs={
        "image": gr.Image(label="Generated Image")
    }
)

# ========== API Integration Node ==========
def call_external_api(data: Dict[str, Any]) -> Dict[str, Any]:
    """Generic API caller with error handling"""
    try:
        response = requests.post(
            "https://api.example.com/v1/process",
            json=data,
            headers={"Authorization": f"Bearer {API_KEYS.get('OPENAI')}"},
            timeout=30
        )
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return {"error": str(e)}

api_integrator = FnNode(
    fn=call_external_api,
    inputs={
        "api_data": gr.JSON(label="API Payload")
    },
    outputs={
        "api_response": gr.JSON(label="API Results")
    }
)

# ========== Output Formatter Node ==========
def format_output(llm_response: str, image: Any, api_data: Dict) -> Dict[str, Any]:
    """Create unified output format"""
    return {
        "text_response": llm_response,
        "visual_response": image,
        "api_data": api_data,
        "status": "success"
    }

output_formatter = FnNode(
    fn=format_output,
    inputs={
        "llm_response": gr.Textbox(),
        "image": gr.Image(),
        "api_data": gr.JSON()
    },
    outputs={
        "final_output": gr.JSON(label="Final Output")
    }
)

# ========== Create and Connect Workflow ==========
workflow = Graph(
    name="Global Integration Platform",
    nodes=[
        input_processor,
        llm_processor,
        image_generator,
        api_integrator,
        output_formatter
    ],
    connections=[
        (input_processor.outputs["processed_data"], llm_processor.inputs["prompt"]),
        (input_processor.outputs["processed_data"], image_generator.inputs["prompt"]),
        (input_processor.outputs["processed_data"], api_integrator.inputs["api_data"]),
        (llm_processor.outputs["response"], output_formatter.inputs["llm_response"]),
        (image_generator.outputs["image"], output_formatter.inputs["image"]),
        (api_integrator.outputs["api_response"], output_formatter.inputs["api_data"])
    ]
)

# ========== Launch Application ==========
if __name__ == "__main__":
    workflow.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        auth=("admin", os.getenv("APP_PASSWORD")),
        favicon_path="https://example.com/favicon.ico"
    )