"""Shared camera-frame analysis for Pipeline and Realtime runtimes."""

from __future__ import annotations

import asyncio
import base64
from dataclasses import dataclass
from io import BytesIO
from typing import Any

from models import (
    VISION_PROMPT_MODES,
    AssistantConfig,
    RuntimeModelResource,
    VisionPromptMode,
)
from openai import AsyncOpenAI
from PIL import Image
from pipecat.frames.frames import UserImageRawFrame


VISION_TOOL_NAME = "fetch_user_image"
VISION_SYSTEM_HINT = (
    "当前阶段打开了视觉理解。用户询问当前画面、摄像头里有什么、人物/物品/"
    "环境状态或需要你看一眼时，调用 fetch_user_image 获取当前视频帧，"
    "再基于画面回答。"
)
VISION_ANALYSIS_SYSTEM_PROMPT = (
    "你是一个视觉理解模型。请只根据图片内容和用户问题给出准确、简洁的中文观察结果。"
    "如果画面不足以判断，请明确说明不确定。"
)


@dataclass(frozen=True)
class VisionInputResolution:
    """The active model path for one explicit user image turn."""

    native: bool
    config: AssistantConfig


def _require(value: str, label: str) -> str:
    if value:
        return value
    raise ValueError(f"缺少模型资源配置: {label}")


def normalize_vision_prompt_mode(value: object) -> VisionPromptMode:
    """Return a supported vision prompt mode."""

    mode = str(value or "runtime")
    if mode == "fixed":
        return "fixed"
    if mode == "guided":
        return "guided"
    return "runtime"


def _compact_prompt(prompt: str, *, limit: int = 4000) -> str:
    return " ".join(prompt.split())[:limit]


def vision_system_hint(
    mode: VisionPromptMode = "runtime",
    configured_prompt: str = "",
) -> str:
    """Tell the conversation model how the camera tool forms its question."""

    prompt_mode = normalize_vision_prompt_mode(mode)
    prompt = _compact_prompt(configured_prompt)
    if prompt_mode == "fixed":
        return (
            f"{VISION_SYSTEM_HINT}\n"
            f"该工具始终使用固定视觉问题“{prompt}”，调用时不需要生成视觉问题。"
        )
    if prompt_mode == "guided":
        return (
            f"{VISION_SYSTEM_HINT}\n"
            "调用时请结合当前对话生成具体视觉问题，并遵守以下固定引导："
            f"{prompt}"
        )
    return (
        f"{VISION_SYSTEM_HINT}\n"
        "调用时请结合当前对话，生成一个具体、可由当前画面回答的视觉问题。"
    )


def vision_tool_schema_parts(
    mode: VisionPromptMode = "runtime",
    configured_prompt: str = "",
) -> tuple[str, dict[str, Any], list[str]]:
    """Build the tool contract for one configured visual-question mode."""

    base = (
        "获取用户当前摄像头画面。当用户询问当前画面、看到了什么、"
        "人/物品/环境状态或需要视觉判断时调用。"
    )
    prompt_mode = normalize_vision_prompt_mode(mode)
    prompt = _compact_prompt(configured_prompt, limit=600)
    if prompt_mode == "fixed":
        return (
            f"{base} 每次都检查同一个固定问题：{prompt}",
            {},
            [],
        )

    question_description = "结合当前对话生成的、关于当前视频画面的具体问题。"
    description = base
    if prompt_mode == "guided":
        description = f"{base} 生成问题时必须遵守固定引导：{prompt}"
        question_description = f"结合当前对话生成具体视觉问题，并遵守固定引导：{prompt}"
    return (
        description,
        {"question": {"type": "string", "description": question_description}},
        ["question"],
    )


def resolve_vision_question(
    mode: VisionPromptMode,
    configured_prompt: str,
    arguments: dict[str, Any] | None,
) -> str:
    """Resolve the exact question sent with the captured image."""

    prompt_mode = normalize_vision_prompt_mode(mode)
    if prompt_mode == "fixed":
        question = configured_prompt.strip()
    else:
        question = str((arguments or {}).get("question") or "").strip()
    if not question:
        raise ValueError(
            "固定视觉提示词为空"
            if prompt_mode == "fixed"
            else "视觉问题为空"
        )
    return question


def image_jpeg_bytes(frame: UserImageRawFrame) -> bytes:
    """Encode one Pipecat camera frame as a storage- and model-ready JPEG."""
    if not frame.format:
        raise ValueError("摄像头图片帧缺少 format,无法编码给视觉模型")
    buffer = BytesIO()
    Image.frombytes(frame.format, frame.size, frame.image).save(
        buffer,
        format="JPEG",
        quality=85,
    )
    return buffer.getvalue()


def image_data_uri(frame: UserImageRawFrame) -> str:
    """Encode one Pipecat camera frame for a vision chat-completion request."""
    encoded = base64.b64encode(image_jpeg_bytes(frame)).decode("utf-8")
    return f"data:image/jpeg;base64,{encoded}"


def image_frame_from_jpeg(data: bytes) -> UserImageRawFrame:
    """Decode one normalized uploaded JPEG into Pipecat's raw image frame."""

    try:
        image = Image.open(BytesIO(data)).convert("RGB")
        image.load()
    except (OSError, ValueError) as exc:
        raise ValueError("上传图片无法解码") from exc
    return UserImageRawFrame(
        image=image.tobytes(),
        size=image.size,
        format="RGB",
    )


async def analyze_image_with_vision_model(
    cfg: AssistantConfig,
    frame: UserImageRawFrame,
    question: str,
) -> str:
    """Analyze one frame with the configured independent vision LLM."""
    if cfg.vision_llm_interface_type not in {"openai-llm", "dashscope-llm"}:
        raise ValueError(
            f"不支持的视觉 LLM 接口类型: {cfg.vision_llm_interface_type}"
        )

    data_uri = await asyncio.to_thread(image_data_uri, frame)
    extra_body = cfg.vision_llm_values.get("extraBody")
    extra = {"extra_body": extra_body} if isinstance(extra_body, dict) else {}
    client = AsyncOpenAI(
        api_key=_require(cfg.vision_llm_api_key, "Vision LLM apiKey"),
        base_url=_require(cfg.vision_llm_base_url, "Vision LLM apiUrl"),
    )
    try:
        response = await client.chat.completions.create(
            model=_require(cfg.vision_model, "Vision LLM modelId"),
            messages=[
                {
                    "role": "system",
                    "content": VISION_ANALYSIS_SYSTEM_PROMPT,
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": question},
                        {"type": "image_url", "image_url": {"url": data_uri}},
                    ],
                },
            ],
            **extra,
        )
    finally:
        await client.close()

    content = response.choices[0].message.content if response.choices else ""
    if isinstance(content, str):
        return content.strip()
    return str(content or "").strip()


def config_with_vision_resource(
    cfg: AssistantConfig,
    resource: RuntimeModelResource,
) -> AssistantConfig:
    """Create the config view used to analyze one Workflow camera frame."""
    if resource.capability != "LLM":
        raise ValueError(f"视觉模型资源能力无效:{resource.capability}")
    if not resource.support_image_input:
        raise ValueError(f"视觉模型不支持图片输入:{resource.id}")

    result = cfg.model_copy(deep=True)
    values = resource.values or {}
    secrets = resource.secrets or {}
    result.vision_model_resource_id = resource.id
    result.vision_model = str(values.get("modelId") or "")
    result.vision_llm_interface_type = resource.interface_type
    result.vision_llm_values = values
    result.vision_llm_secrets = secrets
    result.vision_llm_support_image_input = resource.support_image_input
    result.vision_llm_api_key = str(secrets.get("apiKey") or "")
    result.vision_llm_base_url = str(values.get("apiUrl") or "")
    return result


def config_with_main_llm_as_vision(cfg: AssistantConfig) -> AssistantConfig:
    """Legacy fallback when a Workflow LLM is not present in the resource map."""
    if not cfg.llm_support_image_input:
        raise ValueError("当前大语言模型不支持图片输入")

    result = cfg.model_copy(deep=True)
    result.vision_model_resource_id = None
    result.vision_model = cfg.model
    result.vision_llm_interface_type = cfg.llm_interface_type
    result.vision_llm_values = dict(cfg.llm_values)
    result.vision_llm_secrets = dict(cfg.llm_secrets)
    result.vision_llm_support_image_input = cfg.llm_support_image_input
    result.vision_llm_api_key = cfg.llm_api_key
    result.vision_llm_base_url = cfg.llm_base_url
    return result


def resolve_vision_input(
    cfg: AssistantConfig,
    workflow_scope: dict[str, Any] | None = None,
) -> VisionInputResolution:
    """Resolve native versus auxiliary vision using the active runtime scope."""

    if cfg.type != "workflow":
        if not cfg.vision_enabled:
            raise ValueError("当前助手未开启视觉能力")
        if cfg.vision_model_resource_id:
            if not cfg.vision_llm_support_image_input:
                raise ValueError("当前视觉模型不支持图片输入")
            return VisionInputResolution(native=False, config=cfg)
        if not cfg.llm_support_image_input:
            raise ValueError("当前大语言模型不支持图片输入")
        return VisionInputResolution(native=True, config=cfg)

    scope = workflow_scope or {}
    if not scope.get("enabled"):
        raise ValueError("当前 Workflow Agent 节点未启用视觉能力")
    vision_resource_id = str(scope.get("vision_model_resource_id") or "")
    llm_resource_id = str(scope.get("llm_resource_id") or "")
    resource_id = vision_resource_id or llm_resource_id
    resource = cfg.workflow_model_resources.get(resource_id)
    if resource is not None:
        if not resource.support_image_input:
            raise ValueError(f"当前视觉模型不支持图片输入: {resource.id}")
        result = config_with_vision_resource(cfg, resource)
    elif not vision_resource_id:
        result = config_with_main_llm_as_vision(cfg)
    else:
        raise ValueError(f"视觉模型资源未加载: {vision_resource_id}")

    result.vision_prompt_mode = normalize_vision_prompt_mode(
        scope.get("vision_prompt_mode")
    )
    result.vision_prompt = str(scope.get("vision_prompt") or "").strip()
    return VisionInputResolution(
        native=not bool(vision_resource_id),
        config=result,
    )


async def build_image_user_message(
    cfg: AssistantConfig,
    image_bytes: bytes,
    prompt_text: str,
    *,
    workflow_scope: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Build the same model-facing user message used by an explicit photo turn."""

    question = prompt_text.strip() or "请根据用户刚提交的图片进行回复。"
    frame = await asyncio.to_thread(image_frame_from_jpeg, image_bytes)
    resolution = resolve_vision_input(cfg, workflow_scope)
    if resolution.native:
        data_uri = await asyncio.to_thread(image_data_uri, frame)
        return {
            "role": "user",
            "content": [
                {"type": "text", "text": question},
                {"type": "image_url", "image_url": {"url": data_uri}},
            ],
        }

    observation = await analyze_image_with_vision_model(
        resolution.config,
        frame,
        question,
    )
    return {
        "role": "user",
        "content": (
            f"{question}\n\n"
            "[视觉模型对用户刚提交图片的观察]\n"
            f"{observation or '视觉模型没有返回有效观察结果。'}"
        ),
    }