import base64
import io
import warnings

# Gradio still references Starlette's deprecated HTTP_422_UNPROCESSABLE_ENTITY.
warnings.filterwarnings(
    "ignore",
    message="'HTTP_422_UNPROCESSABLE_ENTITY' is deprecated",
    category=DeprecationWarning,
)
# Stale installs or old Space builds may still import duckduckgo_search; we use Brave only.
warnings.filterwarnings(
    "ignore",
    message=r".*`duckduckgo_search`.*renamed.*`ddgs`.*",
    category=RuntimeWarning,
)

import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image

from agent import run_agent
from agent.orchestrator import strip_qwen_thinking

MODEL = "Qwen/Qwen3-VL-30B-A3B-Thinking"

DEFAULT_SYSTEM = (
    "You are a helpful, multimodal AI assistant. When an image is sent, simply transcribe it without analysis, unless the user asks for analysis." \
    "When you need up-to-date information from the internet, use the search_web tool. Once you have a complete answer, call final_output with the full answer." \
    "If a task is impossible or unsafe, call abort with a brief reason."
)


def _image_to_data_url(image_path: str, max_side: int = 1120, quality: int = 85) -> str:
    """Resize and base64-encode a local image so the HF payload stays under the limit."""
    with Image.open(image_path) as img:
        img = img.convert("RGB")
        w, h = img.size
        if max(w, h) > max_side:
            scale = max_side / max(w, h)
            img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=quality)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def _normalize_content_for_api(content):
    """
    Gradio multimodal history uses {'type': 'file', 'file': FileData}; the HF
    chat API expects {'type': 'image_url', 'image_url': {'url': ...}}.
    """
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return content
    out: list = []
    for part in content:
        if isinstance(part, str):
            out.append({"type": "text", "text": part})
            continue
        if not isinstance(part, dict):
            continue
        ptype = part.get("type")
        if ptype == "text":
            out.append({"type": "text", "text": part.get("text", "")})
        elif ptype == "file":
            fd = part.get("file")
            if isinstance(fd, dict) and fd.get("path"):
                out.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": _image_to_data_url(fd["path"])},
                    }
                )
        elif ptype == "image_url":
            out.append(part)
    if not out:
        return ""
    if len(out) == 1 and out[0].get("type") == "text":
        return out[0].get("text", "")
    return out


def _strip_thinking_from_normalized(content):
    """Remove Qwen thinking tags from text Gradio stored on assistant turns."""
    if isinstance(content, str):
        return strip_qwen_thinking(content)
    if isinstance(content, list):
        out = []
        for part in content:
            if isinstance(part, dict) and part.get("type") == "text":
                t = part.get("text", "")
                out.append({**part, "text": strip_qwen_thinking(t)})
            else:
                out.append(part)
        return out
    return content


def _normalize_history_message(msg: dict) -> dict:
    role = msg.get("role")
    if role not in ("user", "assistant", "system"):
        return msg
    content = msg.get("content")
    normalized = _normalize_content_for_api(content)
    if role == "assistant":
        normalized = _strip_thinking_from_normalized(normalized)
    return {**msg, "content": normalized}


def respond(
    message,
    history: list[dict],
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    client = InferenceClient(token=hf_token.token, model=MODEL)

    # multimodal=True sends {"text": str, "files": [path, ...]}
    # Guard against plain strings in case of edge-case history replay
    if isinstance(message, dict):
        text = message.get("text", "")
        files = message.get("files", [])
    else:
        text = message or ""
        files = []

    if files:
        content = []
        if text:
            content.append({"type": "text", "text": text})
        for f in files:
            content.append({"type": "image_url", "image_url": {"url": _image_to_data_url(f)}})
    else:
        content = text

    messages = [{"role": "system", "content": system_message}]
    messages.extend(_normalize_history_message(m) for m in history)
    messages.append({"role": "user", "content": content})

    answer = run_agent(
        messages=messages,
        client=client,
        model=MODEL,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )

    partial = ""
    for i in range(0, len(answer), 8):
        partial += answer[i : i + 8]
        yield partial


chatbot = gr.ChatInterface(
    respond,
    multimodal=True,
    chatbot=gr.Chatbot(height=700),
    additional_inputs=[
        gr.Textbox(value=DEFAULT_SYSTEM, label="System message"),
        gr.Slider(minimum=1, maximum=16384, value=16384, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()


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