import torch
from langchain_huggingface import HuggingFaceEndpoint
from dotenv import load_dotenv
import streamlit as st
import os
import requests
from PIL import Image
from io import BytesIO

# Load environment variables from .env file
load_dotenv()

# Fetch the Hugging Face API token from the environment
hf_token = os.environ.get("HF_TOKEN")

# Streamlit app title
st.title("Valeey")

# Input text box for user prompt
input_text = st.text_input("Enter your prompt:")

# Dropdown for user to choose the type of generation
generation_type = st.selectbox(
    "What would you like to generate?",
    options=["Text", "Image", "Video"],
    index=0
)

# Sliders for adjustable parameters
if generation_type == "Text":
    max_tokens = st.slider("Max Tokens", min_value=100, max_value=5000, value=1000, step=100)
    temperature = st.slider("Temperature", min_value=0.1, max_value=1.0, value=0.7, step=0.1)
elif generation_type == "Image":
    image_height = st.slider("Image Height", min_value=256, max_value=1024, value=576, step=128)
    image_width = st.slider("Image Width", min_value=256, max_value=1024, value=1024, step=128)
    num_inference_steps = st.slider("Number of Inference Steps", min_value=10, max_value=100, value=50, step=10)
    guidance_scale = st.slider("Guidance Scale", min_value=1.0, max_value=20.0, value=7.5, step=0.5)
elif generation_type == "Video":
    video_height = st.slider("Video Height", min_value=256, max_value=1024, value=576, step=128)
    video_width = st.slider("Video Width", min_value=256, max_value=1024, value=1024, step=128)
    num_frames = st.slider("Number of Frames", min_value=10, max_value=200, value=129, step=10)
    num_inference_steps = st.slider("Number of Inference Steps", min_value=10, max_value=100, value=50, step=10)
    guidance_scale = st.slider("Guidance Scale", min_value=1.0, max_value=20.0, value=7.5, step=0.5)

# Generate button
if st.button("Generate"):
    if input_text:
        if generation_type == "Text":
            # Initialize the HuggingFaceEndpoint LLM for chatbot
            llm = HuggingFaceEndpoint(
                repo_id="mistralai/Mistral-7B-Instruct-v0.3",
                max_new_tokens=max_tokens,
                temperature=temperature,
                huggingfacehub_api_token=hf_token
            )
            
            # Generate the answer using the LLM
            st.write("### Generating Text...")
            try:
                answer = llm.invoke(input_text)
                st.write("### Generated Text:")
                st.write(answer)
            except Exception as e:
                st.error(f"Failed to generate text: {e}")

        elif generation_type == "Image":
            # Function to generate an image using Hugging Face Inference API (Stable Diffusion)
            def generate_image(prompt, hf_token, height, width, num_inference_steps, guidance_scale):
                API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
                headers = {"Authorization": f"Bearer {hf_token}"}
                payload = {
                    "inputs": prompt,
                    "parameters": {
                        "height": height,
                        "width": width,
                        "num_inference_steps": num_inference_steps,
                        "guidance_scale": guidance_scale
                    }
                }
                try:
                    response = requests.post(API_URL, headers=headers, json=payload)
                    response.raise_for_status()
                    image = Image.open(BytesIO(response.content))
                    return image
                except requests.exceptions.RequestException as e:
                    st.error(f"Failed to generate image: {e}")
                    return None

            # Generate an image using the Hugging Face Inference API
            st.write("### Generating Image...")
            image = generate_image(input_text, hf_token, image_height, image_width, num_inference_steps, guidance_scale)
            
            if image:
                # Display the generated image
                st.image(image, caption="Generated Image")

        elif generation_type == "Video":
            # Function to generate a video using HunyuanVideo
            def generate_video(prompt, hf_token, height, width, num_frames, num_inference_steps, guidance_scale):
                API_URL = "https://api-inference.huggingface.co/models/wangfuyun/AnimateLCM"
                headers = {"Authorization": f"Bearer {hf_token}"}
                payload = {
                    "inputs": prompt,
                    "parameters": {
                        "height": height,
                        "width": width,
                        "num_frames": num_frames,
                        "num_inference_steps": num_inference_steps,
                        "guidance_scale": guidance_scale
                    }
                }
                try:
                    response = requests.post(API_URL, headers=headers, json=payload)
                    response.raise_for_status()
                    video_path = "generated_video.mp4"
                    with open(video_path, "wb") as f:
                        f.write(response.content)
                    return video_path
                except requests.exceptions.RequestException as e:
                    st.error(f"Failed to generate video: {e}")
                    return None

            # Generate video using HunyuanVideo
            st.write("### Generating Video...")
            video_path = generate_video(input_text, hf_token, video_height, video_width, num_frames, num_inference_steps, guidance_scale)
            
            if video_path and os.path.exists(video_path):
                # Display the generated video
                st.write("### Generated Video:")
                st.video(video_path)
            else:
                st.error("Failed to generate or locate the video file.")
    else:
        st.warning("Please enter a prompt before generating.")