import gradio as gr
from PIL import Image
from io import BytesIO
import base64
import requests
import os
import random
import torch
import subprocess
import numpy as np
import cv2
from transformers import AutoProcessor, AutoModelForCausalLM
from diffusers import DiffusionPipeline
from datetime import datetime
from mistralai import Mistral
from theme import theme
from fastapi import FastAPI

app = FastAPI()



api_key = os.getenv("MISTRAL_API_KEY")
Mistralclient = Mistral(api_key=api_key)

def flip_image(x):
    return np.fliplr(x)

def encode_image(image_path):
    """Encode the image to base64."""
    try:
        # Open the image file
        image = Image.open(image_path).convert("RGB")

        # Resize the image to a height of 512 while maintaining the aspect ratio
        base_height = 512
        h_percent = (base_height / float(image.size[1]))
        w_size = int((float(image.size[0]) * float(h_percent)))
        image = image.resize((w_size, base_height), Image.LANCZOS)

        # Convert the image to a byte stream
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return img_str
    except FileNotFoundError:
        print(f"Error: The file {image_path} was not found.")
        return None
    except Exception as e:  # Add generic exception handling
        print(f"Error: {e}")
        return None

def feifeichat(image):
    try:
        model = "pixtral-large-2411"
        # Define the messages for the chat
        base64_image = encode_image(image)
        messages = [{
            "role":
            "user",
            "content": [
                {
                    "type": "text",
                    "text": "Please provide a detailed description of this photo"
                },
                {
                    "type": "image_url",
                    "image_url": f"data:image/jpeg;base64,{base64_image}" 
                },
            ],
            "stream": False,
        }]
    
        partial_message = ""
        for chunk in Mistralclient.chat.stream(model=model, messages=messages):
            if chunk.data.choices[0].delta.content is not None:
                partial_message = partial_message + chunk.data.choices[
                    0].delta.content
                yield partial_message
    except Exception as e:  # Add common exception handling
        print(f"Error: {e}")
        return "Please upload a photo"


with gr.Blocks(theme=theme, elem_id="app-container") as app:
    gr.Markdown("Image To Flux Prompt")
    with gr.Tab(label="Image To Prompt"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input Picture",height=320,type="filepath")
                submit_btn = gr.Button(value="Submit", variant='primary')
            with gr.Column():
                output_text = gr.Textbox(label="Flux Prompt", show_copy_button = True)
                clr_button =gr.Button("Clear",variant="primary", elem_id="clear_button")
                clr_button.click(lambda: gr.Textbox(value=""), None, output_text)

        submit_btn.click(feifeichat, [input_img], [output_text])
        
    
        
if __name__ == "__main__":
    app.launch()