import streamlit as st
from daggr import FnNode, Graph
import gradio as gr
from typing import Dict, Any
import os
import time
from datetime import datetime
from utils import validate_metadata, log_workflow

# Set page config
st.set_page_config(
    page_title="Global Integration Platform",
    page_icon="🌐",
    layout="wide"
)

# Add anycoder link
st.sidebar.markdown("[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)")

# Initialize session state
if 'workflow_results' not in st.session_state:
    st.session_state.workflow_results = None

# ========== Nodes Definition ==========
def preprocess_inputs(user_input: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
    """Clean and validate inputs with metadata enrichment"""
    if not user_input.strip():
        raise ValueError("Input cannot be empty")
    return {
        "cleaned_input": user_input.strip(),
        "timestamp": metadata.get("timestamp", datetime.now().isoformat()),
        "source": metadata.get("source", "web")
    }

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

def llm_wrapper(prompt: str, temperature: float = 0.7) -> str:
    """Wrapper for LLM processing with error handling"""
    try:
        # Simulate processing time
        time.sleep(1)
        return f"LLM Response to: {prompt}"
    except Exception as e:
        return f"Error in LLM processing: {str(e)}"

llm_processor = FnNode(
    fn=llm_wrapper,
    inputs={
        "prompt": gr.Textbox(label="LLM Prompt"),
        "temperature": gr.Slider(0, 1, value=0.7)
    },
    outputs={
        "response": gr.Textbox(label="LLM Response")
    }
)

def image_gen_wrapper(prompt: str, negative_prompt: str = "", steps: int = 30) -> Any:
    """Wrapper for image generation"""
    try:
        # Placeholder for actual image generation
        return "https://via.placeholder.com/512?text=Generated+Image"
    except Exception as e:
        return f"Error in image generation: {str(e)}"

image_generator = FnNode(
    fn=image_gen_wrapper,
    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")
    }
)

def call_external_api(data: Dict[str, Any]) -> Dict[str, Any]:
    """Generic API caller with error handling"""
    try:
        # Simulate API call
        time.sleep(0.5)
        return {
            "status": "success",
            "data": {
                "input": data,
                "processed": True,
                "timestamp": datetime.now().isoformat()
            }
        }
    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")
    }
)

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",
        "timestamp": datetime.now().isoformat()
    }

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 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"])
    ]
)

# ========== Streamlit UI ==========
st.title("🌐 Global Integration Platform")
st.markdown("""
This application integrates multiple components into a cohesive workflow including:
- Input processing
- LLM processing
- Image generation
- External API integration
""")

with st.form("workflow_form"):
    user_input = st.text_area("Enter your input:", height=150)
    metadata = st.text_input("Additional metadata (JSON):", value='{"source": "streamlit"}')
    temperature = st.slider("LLM Temperature:", 0.0, 1.0, 0.7)
    steps = st.slider("Image Generation Steps:", 10, 50, 30)
    
    submitted = st.form_submit_button("Run Workflow")
    
    if submitted:
        with st.spinner("Processing workflow..."):
            try:
                # Prepare inputs
                inputs = {
                    "user_input": user_input,
                    "metadata": validate_metadata(metadata),
                    "temperature": temperature,
                    "steps": steps
                }
                
                # Execute workflow
                results = workflow.run(inputs)
                st.session_state.workflow_results = results
                log_workflow(results)
                
                st.success("Workflow completed successfully!")
                
            except Exception as e:
                st.error(f"Error in workflow execution: {str(e)}")

# Display results if available
if st.session_state.workflow_results:
    st.subheader("Workflow Results")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### Text Response")
        st.write(st.session_state.workflow_results['final_output']['text_response'])
        
        st.markdown("### API Response")
        st.json(st.session_state.workflow_results['final_output']['api_data'])
    
    with col2:
        st.markdown("### Generated Image")
        st.image(st.session_state.workflow_results['final_output']['visual_response'])
        
        st.markdown("### Full Output")
        st.json(st.session_state.workflow_results['final_output'])