import os
import gradio as gr
from typing import Callable
import base64
from openai import OpenAI

def get_fn(model_name: str, **model_kwargs):
    """Create a chat function that uses the OpenAI-compatible endpoint."""   

    OPENAI_API_KEY = "-"
    client = OpenAI(
    base_url=" http://192.222.58.60:8000/v1", 
    api_key="tela",
    )

    def predict(
        message: str,
        history,
        system_prompt: str,
        temperature: float,
        top_p: float,
        max_tokens: int,
        
    ):
        try:
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            for user_msg, assistant_msg in history:
                messages.append({"role": "user", "content": user_msg})
                messages.append({"role": "assistant", "content": assistant_msg})
            messages.append({"role": "user", "content": message})

            response = client.chat.completions.create(
                model=model_name,
                messages=messages,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_new_tokens,
                n=1,
                stream=True, 
                response_format={"type": "text"},
            )

            response_text = ""
            for chunk in response:
                chunk_message = chunk.choices[0].delta.content
                if chunk_message:
                    response_text += chunk_message
                    yield assistant_message.strip()
        except Exception as e:
            print(f"Error during generation: {str(e)}")
            yield f"An error occurred: {str(e)}"

    return predict

def get_image_base64(url: str, ext: str):
    with open(url, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    return "data:image/" + ext + ";base64," + encoded_string

def handle_user_msg(message: str):
    if isinstance(message, str):
        return message
    elif isinstance(message, dict):
        if message.get("files"):
            ext = os.path.splitext(message["files"][-1])[1].strip(".")
            if ext.lower() in ["png", "jpg", "jpeg", "gif", "pdf"]:
                encoded_str = get_image_base64(message["files"][-1], ext)
            else:
                raise NotImplementedError(f"Not supported file type {ext}")
            content = [
                {"type": "text", "text": message.get("text", "")},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": encoded_str,
                    }
                },
            ]
        else:
            content = message.get("text", "")
        return content
    else:
        raise NotImplementedError

def get_model_path(name: str = None, model_path: str = None) -> str:
    """Get the model name to use with the endpoint."""
    if model_path:
        return model_path
    if name:
        return name
    raise ValueError("Either name or model_path must be provided")

def registry(name: str = None, model_path: str = None, **kwargs):
    """Create a Gradio ChatInterface."""
    model_name = get_model_path(name, model_path)
    fn = get_fn(model_name, **kwargs)

    interface = gr.ChatInterface(
        fn=fn,
        additional_inputs_accordion=gr.Accordion("⚙️ Parameters", open=False),
        additional_inputs=[
            gr.Textbox(
                "You are a helpful AI assistant.",
                label="System prompt"
            ),
            gr.Slider(0, 1, 0.7, label="Temperature"),
            gr.Slider(128, 4096, 1024, label="Max new tokens"),
            gr.Slider(0, 1, 0.95, label="Top P sampling"),
        ],
    )

    return interface