import base64
import json
import mimetypes
import os

import gradio as gr
from huggingface_hub import get_token
from openai import OpenAI

MODEL = "deepseek-ai/DeepSeek-V4.1-Flash:novita"


def _to_data_url(image) -> str:
    """Workflow image ports arrive as FileData-like dicts, e.g.
    {"path": "/tmp/...", "url": "http://localhost:7860/gradio_api/file=..."}.
    The router can't reach local/localhost URLs, so always read the local
    file and send a base64 data URL. A plain string is assumed to already
    be a URL or path."""
    if isinstance(image, str):
        if image.startswith(("http://", "https://", "data:")):
            return image
    elif isinstance(image, dict):
        path = image.get("path")
        if path and os.path.exists(path):
            mime = mimetypes.guess_type(path)[0] or "image/png"
            with open(path, "rb") as f:
                return f"data:{mime};base64,{base64.b64encode(f.read()).decode()}"
        url = image.get("url") or ""
        if url.startswith(("http://", "https://", "data:")):
            return url
        image = path or url
    raise ValueError(f"Unsupported image value: {str(image)[:200]}")


def _chat(content: list, token: "gr.OAuthToken | None", json_mode: bool = False) -> str:
    """One chat completion against DeepSeek-V4.1-Flash via the HF router.

    `token` is injected by gr.Workflow per request: on Spaces it is the
    visitor's OAuth token; locally it falls back to the host's
    `huggingface_hub login` token. It never appears as a canvas port.
    """
    api_key = token.token if token else get_token()
    client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=api_key)
    kwargs = {"response_format": {"type": "json_object"}} if json_mode else {}
    stream = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": content}],
        stream=True,
        **kwargs,
    )
    output = ""
    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            output += chunk.choices[0].delta.content
    return output


def describe_image(
    image: dict, prompt: str, token: gr.OAuthToken | None = None
) -> str:
    """Describe an uploaded image with DeepSeek-V4.1-Flash."""
    return _chat(
        [
            {"type": "text", "text": prompt},
            {"type": "image_url", "image_url": {"url": _to_data_url(image)}},
        ],
        token,
    )


def extract_details(image: dict, token: gr.OAuthToken | None = None) -> dict:
    """Extract structured details from the image as JSON."""
    text = _chat(
        [
            {
                "type": "text",
                "text": (
                    "Analyze this image and respond with a JSON object with keys: "
                    '"subject" (main subject), "setting" (where it is), '
                    '"colors" (list of dominant colors), "mood" (one word), '
                    '"objects" (list of notable objects).'
                ),
            },
            {"type": "image_url", "image_url": {"url": _to_data_url(image)}},
        ],
        token,
        json_mode=True,
    )
    return json.loads(text)


def write_story(description: str, token: gr.OAuthToken | None = None) -> str:
    """Write a short story based on an image description."""
    return _chat(
        [
            {
                "type": "text",
                "text": (
                    "Write a short, vivid 3-sentence story inspired by this "
                    f"image description:\n\n{description}"
                ),
            }
        ],
        token,
    )


gr.Workflow(
    bind={
        "Describe Image": describe_image,
        "Extract Details": extract_details,
        "Write Story": write_story,
    },
    graph="workflow.json",
).launch()