import os
import tempfile

import gradio as gr
import torch

from PIL import Image

from huggingface_hub import InferenceClient

from ben2 import AutoModel

from smolagents import (
    CodeAgent,
    FinalAnswerTool,
    InferenceClientModel,
    tool,
)
HF_TOKEN = os.getenv("HF_TOKEN")

if HF_TOKEN is None:
    raise ValueError("Add HF_TOKEN to Secrets Hugging Face Space.")

device = "cuda" if torch.cuda.is_available() else "cpu"

client = InferenceClient(
    provider="hf-inference",
    api_key=HF_TOKEN,
)
print("Loading BEN2...")

ben_model = (
    AutoModel
    .from_pretrained("PramaLLC/BEN2")
    .to(device)
)
ben_model.eval()
print("BEN2 loaded.")

model = InferenceClientModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    temperature=0.5,
    max_tokens=2048,
)

@tool
def text_to_image(prompt: str) -> str:
    """Generate a new image.
    Use this tool EVERY TIME the user asks to:
    - generate an image
    - create an image
    - draw
    - render
    - make a picture
    - imagine a scene
    - create art
    Never answer with a text description instead of calling this tool.
    Do not use this tool for editing existing images.
    Args:
        prompt: Detailed description of the image.
    Returns:
        Path to the generated image.
    """
    print("========== TOOL CALLED ==========")
    image = client.text_to_image(
        prompt,
        model="stabilityai/sdxl-turbo",
    )

    output_path = tempfile.mktemp(suffix=".png")

    image.save(output_path)

    return output_path

    
@tool
def change_background_tool(image_path: str) -> str:
    """
    Removes background and places image
    on white background.

    Args:
        image_path: Path to image.

    Returns:
        Path to processed image.
    """

    image = Image.open(image_path).convert("RGB")

    result = ben_model.inference(
        image,
        refine_foreground=True,
    )

    white = Image.new(
        "RGBA",
        result.size,
        (255, 255, 255, 255),
    )

    final = Image.alpha_composite(
        white,
        result,
    ).convert("RGB")

    output_path = tempfile.mktemp(suffix=".png")

    final.save(output_path)

    return output_path
    
agent = CodeAgent(
    tools=[
        text_to_image,
        change_background_tool,
    ],
    model=model,
    stream_outputs=False,
)


def run_agent(prompt, image):

    try:

        # Если пользователь загрузил изображение
        if image is not None:

            if prompt.strip() == "":
                prompt = "Сделай белый фон"

            full_prompt = f"""
Пользователь написал:

{prompt}

Изображение находится по пути:

{image}

Если требуется обработать изображение,
обязательно используй инструмент
change_background_tool.

Не пытайся описывать изображение.
Не придумывай ответ самостоятельно.
"""

        # Если изображения нет
        else:

            full_prompt = f"""
Пользователь написал:

{prompt}

Если пользователь просит создать изображение,
используй инструмент text_to_image.
"""

        result = agent.run(full_prompt)

        # Агент вернул путь к изображению
        if isinstance(result, str):

            if os.path.exists(result):
                return result, "Готово"

        return None, str(result)

    except Exception as e:

        return None, str(e)


########################################################
# Интерфейс
########################################################

demo = gr.Interface(
    fn=run_agent,

    inputs=[
        gr.Textbox(
            label="Запрос",
            lines=3,
            placeholder="Например: сгенерируй красный Ferrari"
        ),

        gr.Image(
            type="filepath",
            label="Изображение (необязательно)"
        ),
    ],

    outputs=[
        gr.Image(label="Результат"),
        gr.Textbox(label="Ответ агента"),
    ],

    title="🤖 AI Image Agent",

    description="""
Загрузите изображение или введите текстовый запрос.

• если изображение загружено —
агент обработает его через BEN2

• если изображения нет —
агент сгенерирует новое изображение.
""",

    allow_flagging="never",
)

########################################################

if __name__ == "__main__":
    demo.launch()