"""
General chat utility functions.
Function 1: Basic chat completion.
Function 2: Image grounding.
"""

import re
from logging import getLogger

from dataclasses import dataclass
from typing import Optional, Literal, Tuple
from PIL import Image, ImageDraw
from openai import OpenAIError

from .model_clients import ClientInfo, DetailedArguments
from .utils import pil_to_base64

Logger = getLogger(__name__)


@dataclass
class ExtraParameters:
    """
    Extra parameters for the chat completion.
    Useful for some models with CoT.
    """

    thinking: Optional[Literal["disabled", "enabled"]] = (
        None  # "thinking" : {"type": "disabled"} for doubao models in volcengine.
    )
    reasoning_effort: Optional[Literal["minimal", "low", "medium", "high"]] = (
        None  # "reasoning_effort" : "medium" for doubao-seed-1-6-lite-251015 or gpt-5
    )


def parse_base64_image(base64_image):
    """
    Parse base64 image to dict with media_type and data for anthropic.
    """
    match = re.search(r"data:(image/[^;]+);base64,", base64_image)
    return {
        "type": "base64",
        "media_type": match.group(1) if match else "image/jpeg",
        "data": base64_image.split(",")[-1] if "," in base64_image else base64_image,
    }


def chat_completion(
    client_info: ClientInfo,
    system_prompt: str,
    user_prompt: str,
    images: list[Image.Image] | None = None,
    temperature: float = 1.0,
    top_p: float = 0.95,
    top_k: int = 40,
    max_tokens: int = 1024,
    unload_after_chat: bool = True,
    extra_parameters: Optional[ExtraParameters] = None,
) -> str:
    """
    Basic chat completion.
    """
    # Basic check of the user's input.
    if images is None:
        payload_images: list = []
    elif len(images) > 4:
        Logger.warning(
            "Warning: The number of images is greater than 4. Only the first 4 images will be used."
        )
        payload_images: list = images[:4]
    else:
        payload_images: list = images

    base64_images: list[str] = [pil_to_base64(image) for image in payload_images]
    orig_args: DetailedArguments = client_info.arguments
    messages: list[dict] = [
        {
            "role": "system",
            "content": system_prompt,
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": user_prompt},
                *[
                    {"type": "image_url", "image_url": {"url": base64_image}}
                    for base64_image in base64_images
                ],
            ],
        },
    ]
    # prepare the arguments
    # Due to ollama issue (Cannot unload model using openai chat api.).
    # We have to generate a single payload here.
    if client_info.client_type == "ollama":
        payload = {
            "model": orig_args.model,
            "messages": [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content": user_prompt,
                    "images": [base64_image for base64_image in base64_images],
                },
            ],
            "options": {
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
                "num_predict": max_tokens,
            },
        }
        if extra_parameters:
            # "minimal" thinking effect is not supported by ollama.
            # but "minimal" thinking effect belong to no think is volcengine.
            need_think = extra_parameters.thinking == "enabled"
            payload["think"] = need_think
            if need_think and extra_parameters.reasoning_effort != "minimal":
                payload["think"] = extra_parameters.reasoning_effort
        if unload_after_chat:
            payload["keep_alive"] = "0"

    elif client_info.client_type == "openai":
        payload = {
            "model": orig_args.model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens,
        }
        if extra_parameters:
            if extra_parameters.thinking == "enabled":
                if "seed-1-6" in orig_args.model:
                    payload["extra_body"] = {}
                    payload["extra_body"]["thinking"] = {
                        "type": extra_parameters.thinking
                    }
                payload["reasoning_effort"] = extra_parameters.reasoning_effort

    elif client_info.client_type == "openai-responses":
        payload = {
            "model": orig_args.model,
            "input": [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": user_prompt},
                        *[
                            {"type": "input_image", "image_url": base64_image}
                            for base64_image in base64_images
                        ],
                    ],
                },
            ],
            "temperature": temperature,
            "top_p": top_p,
            "max_output_tokens": max_tokens,
        }
        if extra_parameters:
            model: str = orig_args.model
            if "seed-1-6" in model:
                payload["extra_body"] = {}
                payload["extra_body"]["thinking"] = {"type": extra_parameters.thinking}
            else:
                # volcengine does not support thinking effort in responses api
                payload["reasoning"] = {}
                payload["reasoning"]["effort"] = extra_parameters.reasoning_effort

    elif client_info.client_type == "mistral":
        payload = {
            "model": orig_args.model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "max_tokens": max_tokens,
        }

    elif client_info.client_type == "anthropic":
        anthropic_messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_prompt},
                    *[
                        {
                            "type": "image",
                            "source": parse_base64_image(base64_image),
                        }
                        for base64_image in base64_images
                    ],
                ],
            },
        ]
        payload = {
            "model": orig_args.model,
            "system": [{"text": system_prompt, "type": "text"}],
            "messages": anthropic_messages,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "max_tokens": max_tokens,
            "thinking": {"type": "disabled"},
        }
        if extra_parameters and extra_parameters.thinking == "enabled":
            payload["thinking"] = {"type": extra_parameters.thinking}
            if (
                extra_parameters.reasoning_effort
                and extra_parameters.reasoning_effort != "minimal"
            ):
                effort_token = {
                    "low": 1024,
                    "medium": 2048,
                    "high": 4096,
                }
                payload["thinking"]["budget_tokens"] = effort_token.get(  # type: ignore
                    extra_parameters.reasoning_effort, 1024
                )
            else:
                payload["thinking"] = {"type": "disabled"}

    else:
        raise ValueError("The client type is not supported.")

    # Call the API
    try:
        response = client_info.chat_func(**payload)
    except OpenAIError as e:
        raise RuntimeError(f"Error in openai chat completion: {e}") from e
    except Exception as e:
        raise RuntimeError(f"Unexpected error in chat completion: {e}") from e

    # Parse the response
    if client_info.client_type in ["openai", "mistral"]:  # Chat Completion API
        result: str = response.choices[0].message.content
    elif client_info.client_type == "openai-responses":  # Responses API
        result: str = response.output_text
    elif client_info.client_type == "ollama":
        result: str = response.message.content
    elif client_info.client_type == "anthropic":
        result: str = response.content[0].text
    else:
        raise ValueError("The message type is not supported.")
    return result


def grounding(
    client_info: ClientInfo,
    item: Optional[str],
    image: Image.Image,
    thinking: Literal["disabled", "enabled"] = "disabled",
    mode: Literal["minimal", "low", "medium", "high"] = "minimal",
    unload_after_chat: bool = True,
) -> Tuple[list[list[int]], Image.Image]:
    """
    Image grounding with item.
    """
    # Basic check of the user's input.
    if item is None:
        raise ValueError("The item is None")

    # prepare the prompt
    system_prompt = "You are a professional image grounding assistant."
    user_prompt = (
        f"Please locate the item '{item}' in the image accurately. "
        + "Response in coordinate of the bounding box. "
        + "The format is <bbox>x_min y_min x_max y_max</bbox> in percentage(0-1000). "
        + "If there are multiple items, please list all bounding boxes. "
    )

    # prepare the extra parameters
    extra_parameters = ExtraParameters(thinking=thinking, reasoning_effort=mode)

    # call the API
    response: str = chat_completion(
        client_info=client_info,
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        images=[image],
        temperature=1.0,
        top_p=0.7,
        max_tokens=4096,
        unload_after_chat=unload_after_chat,
        extra_parameters=extra_parameters,
    )

    # responses: <bbox>0 55 63 250</bbox><bbox>0 300 51 427</bbox>...
    bboxes: list[list[int]] = []
    for bbox in re.findall(r"<bbox>(.*?)</bbox>", response):
        # seed 1.6
        bbox = [int(x) for x in bbox.split()]
        if len(bbox) != 4:
            raise ValueError("BBox is invaild. Please retry again.")
        x_min, y_min, x_max, y_max = tuple(bbox)
        w, h = image.size
        x_min_real = int(x_min * w / 1000)
        y_min_real = int(y_min * h / 1000)
        x_max_real = int(x_max * w / 1000)
        y_max_real = int(y_max * h / 1000)

        bbox_x = x_min_real
        bbox_y = y_min_real
        bbox_w = x_max_real - x_min_real
        bbox_h = y_max_real - y_min_real

        bboxes.append([bbox_x, bbox_y, bbox_w, bbox_h])

    # Create a copy of the image to draw on
    draw_image = image.copy()
    draw = ImageDraw.Draw(draw_image)

    # Define bbox style
    bbox_color = "red"  # Red color for bboxes
    bbox_width = 2  # Line width

    # Draw each bbox
    for bbox in bboxes:
        if len(bbox) == 4:  # Ensure bbox has 4 coordinates [x1, y1, x2, y2]
            x1, y1, x2, y2 = bbox
            # Draw rectangle
            # Pillow uses [x1, x2, x3, x4]
            draw.rectangle(
                [x1, y1, x1 + x2, y1 + y2], outline=bbox_color, width=bbox_width
            )

    return bboxes, draw_image


def captioning(
    client_info: ClientInfo,
    image: Image.Image,
    language: str = "English",
    num_max_sentences: int = 10,
    thinking: Literal["disabled", "enabled"] = "disabled",
    mode: Literal["minimal", "low", "medium", "high"] = "minimal",
    unload_after_chat: bool = True,
) -> str:
    """
    Image captioning.
    """

    # Captioning prompt based on Qwen-Image arxiv paper.
    # https://arxiv.org/pdf/2508.02324
    # Figure 12 in Page 13.
    # Modified by Deepseek-V3.2-Exp to remove unused json format.
    # Better for using thinking mode with "enabled" and "medium" reasoning_effort.
    system_prompt = (
        "As a professional image annotator, "
        + "your primary task is to generate a detailed and natural caption for the input image, "
        + "focusing on authenticity and accuracy without any generalizations. "
        + "Write the caption in descriptive, flowing text, "
        + "avoiding structured formats or rich text elements. "
        + "Enrich the description by including specific object attributes, "
        + "visual relationships between objects, and environmental context "
        + "to provide a comprehensive view. "
        + "If any text is visible in the image, identify it exactly as seen "
        + "and highlight it within the caption using quotation marks, "
        + "without translating or explaining the content. "
        + "Ensure all details are grounded in what is actually present, "
        + "maintaining a truthful representation of the scene. "
        + "This approach guarantees a clear and informative caption "
        + "that captures the essence of the image effectively."
    )
    user_prompt = (
        "Here's the user's input. "
        + f"Please generate a detailed caption for the image in language {language} "
        + f"with maxium {num_max_sentences} sentences."
    )

    # prepare the extra parameters
    extra_parameters = ExtraParameters(thinking=thinking, reasoning_effort=mode)

    # call the API
    response: str = chat_completion(
        client_info=client_info,
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        images=[image],
        temperature=1.0,
        top_p=0.7,
        max_tokens=4096,
        unload_after_chat=unload_after_chat,
        extra_parameters=extra_parameters,
    )

    return response