"""GLM adapter using official ZhipuAI zai-sdk.

Official SDK: zai-sdk (not zhipuai)
Documentation: https://docs.bigmodel.cn/
"""

import base64
import mimetypes
import warnings
from pathlib import Path
from typing import List, Optional

from zai import ZhipuAiClient

from .base import VendorAdapter
from ..models import Response, Usage


class GLMAdapter(VendorAdapter):
    """GLM adapter using official ZhipuAI zai-sdk.

    Supports:
    - Text chat via chat.completions API
    - Vision chat via chat.completions API with image_url content
    - OCR via layout_parsing API for glm-ocr model
    - Optional timeout for requests (if SDK supports)
    """

    DEFAULT_MODEL = "glm-4.5-air"

    # Models that use layout_parsing API (OCR-specific)
    LAYOUT_PARSING_MODELS = {"glm-ocr"}

    def __init__(
        self,
        api_key: str,
        model: Optional[str] = None,
        base_url: Optional[str] = None,
        timeout: Optional[float] = None,
        **kwargs,
    ):
        super().__init__(
            api_key=api_key,
            model=model or self.DEFAULT_MODEL,
            base_url=base_url,
            timeout=timeout,
            **kwargs,
        )

        # Use official zai-sdk client
        # Note: zai-sdk may or may not support timeout parameter
        client_kwargs: dict = {"api_key": self.api_key}
        if self.timeout is not None:
            client_kwargs["timeout"] = self.timeout

        try:
            self.client = ZhipuAiClient(**client_kwargs)
        except TypeError:
            # Older SDK versions may not support the timeout parameter
            client_kwargs.pop("timeout", None)
            self.client = ZhipuAiClient(**client_kwargs)

    def _is_image_file(self, path: Path) -> bool:
        mime_type, _ = mimetypes.guess_type(str(path))
        return bool(mime_type and mime_type.startswith("image/"))

    def _encode_file_as_data_url(self, path: Path) -> str:
        """Encode image or PDF file as base64 data URL."""
        self._check_file_size(path)  # Prevent OOM on large files
        mime_type, _ = mimetypes.guess_type(str(path))
        if not mime_type or (not mime_type.startswith("image/") and mime_type != "application/pdf"):
            raise ValueError(
                f"GLM multimodal/OCR only accepts image or PDF files: {path.name}"
            )

        with open(path, "rb") as f:
            payload = base64.b64encode(f.read()).decode("utf-8")
        return f"data:{mime_type};base64,{payload}"

    def _extract_content(self, completion) -> str:
        """Extract text content from completion response."""
        if not getattr(completion, "choices", None):
            return ""

        message = completion.choices[0].message
        content = getattr(message, "content", None)
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = []
            for item in content:
                if isinstance(item, str):
                    parts.append(item)
                elif isinstance(item, dict):
                    text = item.get("text")
                    if text:
                        parts.append(text)
            return "\n".join(parts)
        return content if content is not None else ""

    def _extract_layout_content(self, response) -> str:
        """Extract text content from layout_parsing response."""
        # layout_parsing returns LayoutParsingResp with md_results field
        if hasattr(response, "md_results") and response.md_results:
            return response.md_results
        if hasattr(response, "text") and response.text:
            return response.text
        if hasattr(response, "content"):
            content = response.content
            if isinstance(content, str):
                return content
        if isinstance(response, dict):
            return response.get("md_results", "") or response.get("text", "") or response.get("content", "")
        return str(response) if response else ""

    def upload_file(self, file_path: Path) -> str:
        path = self._validate_file(file_path)
        raise NotImplementedError(
            f"GLM upload_file() is not supported yet for chat workflows: {path.name}"
        )

    def chat(
        self,
        prompt: str,
        files: Optional[List[Path]] = None,
        **kwargs,
    ) -> Response:
        """Send chat request to GLM using the official zai-sdk.

        For glm-ocr model: uses layout_parsing API
        For other models: uses chat.completions API
        """
        # glm-ocr uses layout_parsing API
        if self.model.lower() in self.LAYOUT_PARSING_MODELS:
            return self._chat_layout_parsing(prompt, files, **kwargs)

        # Other models use chat.completions API
        return self._chat_completions(prompt, files, **kwargs)

    def _chat_completions(
        self,
        prompt: str,
        files: Optional[List[Path]] = None,
        **kwargs,
    ) -> Response:
        """Send chat request using chat.completions API."""
        validated_files = [self._validate_file(Path(file_path)) for file_path in files] if files else []

        if validated_files:
            # Check for unsupported file types (allow images and PDFs)
            supported_extensions = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".pdf"}
            unsupported = [p for p in validated_files if p.suffix.lower() not in supported_extensions]
            if unsupported:
                file_names = ", ".join(path.name for path in unsupported)
                raise ValueError(
                    f"GLM adapter only supports image and PDF files for multimodal chat; "
                    f"unsupported files: {file_names}"
                )

            # Build multimodal content following official zai-sdk format
            content = [{"type": "text", "text": prompt}]
            content.extend(
                {
                    "type": "image_url",
                    "image_url": {"url": self._encode_file_as_data_url(path)},
                }
                for path in validated_files
            )
            messages = [{"role": "user", "content": content}]
        else:
            messages = [{"role": "user", "content": prompt}]

        # Build request kwargs
        request_kwargs = {
            "model": self.model,
            "messages": messages,
        }
        if "temperature" in kwargs:
            request_kwargs["temperature"] = kwargs.pop("temperature")
        if "max_tokens" in kwargs:
            request_kwargs["max_tokens"] = kwargs.pop("max_tokens")
        request_kwargs.update(kwargs)

        completion = self.client.chat.completions.create(**request_kwargs)
        return Response(
            content=self._extract_content(completion),
            usage=Usage.from_openai(getattr(completion, "usage", None)),
            raw=completion,
        )

    def _chat_layout_parsing(
        self,
        prompt: str,
        files: Optional[List[Path]] = None,
        **kwargs,
    ) -> Response:
        """Send request for glm-ocr model.

        Behavior:
        - If prompt is non-empty: uses chat.completions API (supports custom prompt)
        - If prompt is empty: uses layout_parsing API (preserves document layout)

        The layout_parsing API is best for complex document layout preservation,
        while chat.completions is better for structured data extraction with
        custom prompts (e.g., "extract invoice number as JSON").

        Note: layout_parsing API only supports single file (PNG/JPG/PDF).
              chat.completions supports multiple images.
        """
        if not files:
            raise ValueError(
                f"GLM model '{self.model}' requires image/PDF file input. "
                f"Please provide files parameter."
            )

        validated_files = [self._validate_file(Path(file_path)) for file_path in files]

        # Check file types
        supported_types = {".png", ".jpg", ".jpeg", ".pdf"}
        for path in validated_files:
            if path.suffix.lower() not in supported_types:
                raise ValueError(
                    f"GLM OCR model supports PNG, JPG, JPEG, PDF formats only; "
                    f"unsupported file: {path.name}"
                )

        # If prompt is provided, use chat.completions for custom instruction support
        if prompt and prompt.strip():
            return self._chat_completions(prompt, validated_files, **kwargs)

        # Empty prompt: use layout_parsing for layout-preserving OCR
        # layout_parsing only supports single file
        if len(validated_files) > 1:
            raise ValueError(
                f"GLM OCR model with layout_parsing supports single file only; "
                f"got {len(validated_files)} files. "
                f"Use custom prompt to enable chat.completions mode for multi-file support."
            )

        file_path = validated_files[0]
        file_url = self._encode_file_as_data_url(file_path)

        response = self.client.layout_parsing.create(
            model=self.model,
            file=file_url,
        )

        return Response(
            content=self._extract_layout_content(response),
            usage=Usage.from_openai(getattr(response, "usage", None)),
            raw=response,
        )
