#!/usr/bin/env python3
"""
Image description via a generic multimodal VLM.

This module provides:
- a Python API: describe_image(...)
- a CLI: python -m viseeker.image_describe ...

Inputs can be:
- local image file paths
- HTTP/HTTPS URLs
- S3 URLs (s3://bucket/key)

Internally it:
- resolves the input path using PreparedInput (download mode)
- base64-encodes the image
- calls a multimodal chat completion endpoint compatible with OpenAI's Python SDK
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
from typing import Optional

from dotenv import load_dotenv

from ._internal import inputs

# Load environment variables from a .env file if present.
load_dotenv()

try:
    # OpenAI-style client compatible with VLM HTTP endpoint.
    from openai import OpenAI
except ImportError as e:  # pragma: no cover - import error surfaced at runtime
    raise ImportError(
        "The 'openai' package is required for image_describe.\n"
        "Install it with: pip install openai"
    ) from e


DEFAULT_PROMPT = (
    "You are an expert visual analyst. Carefully describe the image in concise, structured text.\n"
    "- Focus on the main objects, their relationships, and overall scene.\n"
    "- Mention important text in the image if any.\n"
    "- Avoid speculation about things that cannot be seen clearly."
)


def _load_image_as_base64(path: str) -> str:
    """
    Read an image file from disk and return a base64-encoded string.
    """
    with open(path, "rb") as f:
        data = f.read()
    return base64.b64encode(data).decode("utf-8")


def _guess_image_mime_type(path: str) -> str:
    """
    Guess image MIME type from file extension.

    Supported mappings (case-insensitive):
    - .jpg, .jpeg -> image/jpeg
    - .png        -> image/png
    - .gif        -> image/gif
    - .webp       -> image/webp
    - .bmp, .dib  -> image/bmp
    - .tiff, .tif -> image/tiff
    - .ico        -> image/ico
    - .icns       -> image/icns
    - .sgi        -> image/sgi
    - .j2c, .j2k, .jp2, .jpc, .jpf, .jpx -> image/jp2
    - .heic       -> image/heic
    - .heif       -> image/heif

    For unknown extensions, fall back to image/jpeg.
    """
    _, ext = os.path.splitext(path)
    ext = ext.lower()
    mapping = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".bmp": "image/bmp",
        ".dib": "image/bmp",
        ".tiff": "image/tiff",
        ".tif": "image/tiff",
        ".ico": "image/ico",
        ".icns": "image/icns",
        ".sgi": "image/sgi",
        ".j2c": "image/jp2",
        ".j2k": "image/jp2",
        ".jp2": "image/jp2",
        ".jpc": "image/jp2",
        ".jpf": "image/jp2",
        ".jpx": "image/jp2",
        ".heic": "image/heic",
        ".heif": "image/heif",
    }
    return mapping.get(ext, "image/jpeg")


def _create_client() -> OpenAI:
    """
    Create an OpenAI-compatible client for the configured VLM endpoint.
    """
    resolved_api_key = os.getenv("VLM_API_KEY")
    if not resolved_api_key:
        raise ValueError("Missing API key: set VLM_API_KEY in environment.")
    resolved_base_url = os.getenv("VLM_BASE_URL")
    if not resolved_base_url:
        raise ValueError("Missing base URL: set VLM_BASE_URL in environment.")
    return OpenAI(base_url=resolved_base_url, api_key=resolved_api_key)


def describe_image(
    input_path: str,
    *,
    prompt: Optional[str] = None,
) -> str:
    """
    Describe an image using a multimodal VLM via an OpenAI-compatible client.

    Args:
        input_path: Local path, HTTP/HTTPS URL, or S3 URL to the image.
        prompt: Optional custom prompt. When None, a built-in default prompt is used.

    Returns:
        Description text generated by the VLM (string, JSON-serializable).
    """
    used_prompt = prompt or DEFAULT_PROMPT

    # 模型 ID：必须从环境变量 VLM_MODEL_ID 提供。
    used_model = os.getenv("VLM_MODEL_ID")
    if not used_model:
        raise ValueError("Missing model id: set VLM_MODEL_ID in environment.")

    client = _create_client()

    # Prepare local image file (handles local/HTTP/S3) and detect MIME type.
    with inputs.PreparedInput(input_path, mode="download") as local_path:
        image_b64 = _load_image_as_base64(local_path)
        mime_type = _guess_image_mime_type(local_path)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:{mime_type};base64,{image_b64}",
                    },
                },
                {
                    "type": "text",
                    "text": used_prompt,
                },
            ],
        }
    ]

    try:
        resp = client.chat.completions.create(
            model=used_model,
            messages=messages,
        )
    except Exception as e:
        raise ValueError(f"Image description request failed: {e}") from e

    # Best-effort extraction of response text from OpenAI-style response.
    description_text = ""
    try:
        choice = resp.choices[0]
        content = getattr(choice.message, "content", None)
        if isinstance(content, str):
            description_text = content
        elif isinstance(content, list) and content:
            # Ark/Doubao may return list of content parts, similar to OpenAI multi-modal.
            parts = []
            for part in content:
                if isinstance(part, dict):
                    text = part.get("text")
                    if text:
                        parts.append(text)
            description_text = "\n".join(parts).strip()
        if not description_text and hasattr(choice.message, "content"):
            description_text = str(choice.message.content)
    except Exception:
        description_text = ""

    return description_text


def _build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="image_describe",
        description="Describe an image using a multimodal VLM via an OpenAI-compatible client.",
    )
    parser.add_argument(
        "input_path",
        help="Image path/URL (e.g., ./a.png, https://..., s3://bucket/key).",
    )
    parser.add_argument(
        "--prompt",
        default=None,
        help="Custom prompt to guide the description. Overrides the built-in default prompt.",
    )
    return parser


def main(argv: Optional[list[str]] = None) -> int:
    parser = _build_arg_parser()
    args = parser.parse_args(argv)

    try:
        result = describe_image(
            args.input_path,
            prompt=args.prompt,
        )
        print(json.dumps(result, indent=2, ensure_ascii=False))
        return 0
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
