import gradio as gr
from collections.abc import Generator
from openai import OpenAI
# from gradio.chat_interface import ChatInterface
from chat_interface import ChatInterface

USERNAME = "ahmedheakl"
SPACE_NAME = "arabic-vlm-app"
TITLE = "AIN Arabic VLM"
DESCRIPTION = "Welcome to the AIN Arabic VLM chatbot. The best Arabic-English VLM developed by MBZUAI."
TOP_N_HISTORY = 2
LOGO_PATH = "https://huggingface.co./spaces/ahmedheakl/arabic-vlm-app/resolve/main/logo.jpeg"



def load_chat(
    base_url: str,
    model: str,
    token: str | None = None,
    *,
    system_message: str | None = None,
    **kwargs,
) -> gr.ChatInterface:
    client = OpenAI(api_key=token, base_url=base_url)
    start_message = (
        [{"role": "system", "content": system_message}] if system_message else []
    )

    def open_api_stream(
        message: str, history: list | None
    ) -> Generator[str, None, None]:
        history = history or start_message
        if len(history) > 0 and isinstance(history[0], (list, tuple)):
            history = history[:TOP_N_HISTORY]
            history = ChatInterface._tuples_to_messages(history)
        files = message.get('files', [])
        content = [
            {"type": "text", "text": message.get('text', '')}
        ]
        if files:
            src_path = files[0]
            image_url = f"https://{USERNAME}-{SPACE_NAME}.hf.space/gradio_api/file={src_path}"
            content.append({"type": "image_url", "image_url": {"url": image_url}})
        stream = client.chat.completions.create(
            model=model,
            messages=history + [{"role": "user", "content": content}],
            stream=True,
        )
        response = ""
        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                response += chunk.choices[0].delta.content
                yield response

    return gr.ChatInterface(
        open_api_stream, type="messages", **kwargs
    )

load_chat(
    "https://0f21-5-195-0-150.ngrok-free.app/v1",
    model="test",
    token="ollama",
    multimodal=True,
    title=TITLE,
    description=DESCRIPTION,
    theme="ocean",
    examples=[
        {
            "text": "أخبرني ما اسم المبنى الموجود في الصورة والشيء المثير للاهتمام فيه",
            "files": ["https://cdn.mos.cms.futurecdn.net/5HrnHp9ybAqYrtruKAsfkN-1200-80.jpg"],
        },
        {
            "text": "ما هو العلم الموجود في الصورة؟", 
            "files": ["https://mtc.ae/wp-content/uploads/2023/09/Satin-UAE-Flag-UAE-F-B-Blank.jpg"],
        },
        {
            "text": "How many people are there in the image?",
            "files": ["https://i0.wp.com/eatpitapita.com/wp-content/uploads/2020/02/Arab-Muslim-or-Middle-Eastern-Preview.jpg"]
        },
    ],
).queue().launch(allowed_paths=["/static"])