import gradio as gr
import os
import tempfile
from pathlib import Path
import secrets
from PIL import Image
import requests
import json

# API key setup
YOUR_API_TOKEN = os.getenv('YOUR_API_TOKEN')

# API Endpoints (Gerçek API URL'lerinizi buraya ekleyin)
IMAGE_DESCRIPTION_API_ENDPOINT = "https://api.actual-qwen-model.com/v1/image-description"
MATH_RESPONSE_API_ENDPOINT = "https://api.actual-qwen-model.com/v1/math-response"

def process_image(image, should_convert=False):
    # Geçici dizin oluşturma
    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(Path(tempfile.gettempdir()) / "gradio")
    os.makedirs(uploaded_file_dir, exist_ok=True)

    # Geçici dosya adı oluşturma
    name = f"tmp{secrets.token_hex(20)}.jpg"
    filename = os.path.join(uploaded_file_dir, name)

    # İsteğe bağlı olarak resmi dönüştürme
    if should_convert:
        new_img = Image.new('RGB', size=(image.width, image.height), color=(255, 255, 255))
        new_img.paste(image, (0, 0), mask=image)
        image = new_img

    # Resmi kaydetme
    image.save(filename)

    # API isteği hazırlama
    headers = {
        "Authorization": f"Bearer {YOUR_API_TOKEN}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "qwen-vl-max-0809",
        "messages": [
            {
                'role': 'system',
                'content': [{'text': 'You are a helpful assistant.'}]
            },
            {
                'role': 'user',
                'content': [
                    {'image': f'file://{filename}'},
                    {'text': 'Please describe the math-related content in this image, ensuring that any LaTeX formulas are correctly transcribed. Non-mathematical details do not need to be described.'}
                ]
            }
        ]
    }

    try:
        # API isteğini gönderme
        response = requests.post(IMAGE_DESCRIPTION_API_ENDPOINT, headers=headers, json=data)
        response.raise_for_status()
        response_data = response.json()

        # Açıklamayı çıkarma
        description = response_data["output"]["choices"][0]["message"]["content"]

    except Exception as e:
        description = f"Error processing image: {str(e)}"

    finally:
        # Geçici dosyayı silme
        os.remove(filename)

    return description

def get_math_response(image_descriptions, user_question):
    if not image_descriptions:
        return "No image descriptions provided."

    # API isteği hazırlama
    headers = {
        "Authorization": f"Bearer {YOUR_API_TOKEN}",
        "Content-Type": "application/json"
    }
    content = "Image descriptions:\n" + "\n".join(image_descriptions) + f"\n\nUser question: {user_question}"
    data = {
        "model": "qwen2.5-math-72b-instruct",
        "messages": [
            {'role': 'system', 'content': 'You are a helpful math assistant.'},
            {'role': 'user', 'content': content}
        ],
        "result_format": "message",
        "stream": False  # Streaming'i kapattık
    }

    try:
        # API isteğini gönderme
        response = requests.post(MATH_RESPONSE_API_ENDPOINT, headers=headers, json=data)
        response.raise_for_status()
        response_data = response.json()

        # Yanıtı çıkarma
        answer = response_data["output"]["choices"][0]["message"]["content"]

    except Exception as e:
        answer = f"Error generating response: {str(e)}"

    return answer

def math_chat_bot(images, sketchpad, question, chat_history):
    image_descriptions = []

    # Yeni resimleri işleme
    if images is not None:
        for image in images:
            if image:
                img = Image.open(image.name) if hasattr(image, 'name') else image
                description = process_image(img)
                image_descriptions.append(description)

    # Sketchpad varsa işle
    if sketchpad and sketchpad["composite"]:
        sketch_image = sketchpad["composite"]
        sketch_description = process_image(sketch_image, should_convert=True)
        image_descriptions.append(sketch_description)

    # Matematik yanıtını al
    math_response = get_math_response(image_descriptions, question)

    # Sohbet geçmişine ekleme
    if chat_history is None:
        chat_history = []
    chat_history.append((question, math_response))

    return chat_history

css = """
#qwen-md .katex-display { display: inline; }
#qwen-md .katex-display>.katex { display: inline; }
#qwen-md .katex-display>.katex>.katex-html { display: inline; }
"""

# Gradio arayüzünü oluşturma
with gr.Blocks(css=css) as demo:
    gr.HTML("""\
    <p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 60px"/></p>
    <center><font size=8>📖 Qwen2.5-Math Demo</font></center>
    <center><font size=3>This WebUI is based on Qwen2-VL for OCR and Qwen2.5-Math for mathematical reasoning. You can input either images or texts of mathematical or arithmetic problems.</font></center>
    """)

    with gr.Row():
        with gr.Column():
            input_images = gr.File(file_count="multiple", label="Upload Images")
            input_sketchpad = gr.Sketchpad(type="pil", label="Sketch", layers=False)
            input_text = gr.Textbox(label="Input your question")
            with gr.Row():
                clear_btn = gr.ClearButton([input_images, input_sketchpad, input_text])
                submit_btn = gr.Button("Submit", variant="primary")

        with gr.Column():
            chat_output = gr.Chatbot(label="Chat History", elem_id="qwen-md")

    submit_btn.click(
        fn=math_chat_bot,
        inputs=[input_images, input_sketchpad, input_text, chat_output],
        outputs=chat_output
    )

demo.launch()