# -*- coding: utf-8 -*-
"""Message processing utilities for agent communication.

This module handles:
- File and media block processing
- Message content manipulation
- Message validation
"""
import asyncio
import logging
import os
import urllib.parse
from pathlib import Path
from typing import Optional

from agentscope.message import Msg, TextBlock

from ...config import load_config
from .file_handling import download_file_from_base64, download_file_from_url

logger = logging.getLogger(__name__)


async def _process_single_file_block(
    source: dict,
    filename: Optional[str],
) -> Optional[str]:
    """
    Process a single file block and download the file.

    Args:
        source: The source dict containing file information.
        filename: The filename to save.

    Returns:
        The local file path if successful, None otherwise.
    """
    if isinstance(source, dict) and source.get("type") == "base64":
        if "data" in source:
            base64_data = source.get("data", "")
            local_path = await download_file_from_base64(
                base64_data,
                filename,
            )
            logger.debug(
                "Processed base64 file block: %s -> %s",
                filename or "unnamed",
                local_path,
            )
            return local_path

    elif isinstance(source, dict) and source.get("type") == "url":
        url = source.get("url", "")
        if url:
            local_path = await download_file_from_url(
                url,
                filename,
            )
            logger.debug(
                "Processed URL file block: %s -> %s",
                url,
                local_path,
            )
            return local_path

    return None


def _extract_source_and_filename(block: dict, block_type: str):
    """Extract source and filename from a block."""
    if block_type == "file":
        return block.get("source", {}), block.get("filename")

    source = block.get("source", {})
    if not isinstance(source, dict):
        return None, None

    filename = None
    if source.get("type") == "url":
        url = source.get("url", "")
        if url:
            parsed = urllib.parse.urlparse(url)
            filename = os.path.basename(parsed.path) or None

    return source, filename


def _media_type_from_path(path: str) -> str:
    """Infer audio media_type from file path suffix."""
    ext = (os.path.splitext(path)[1] or "").lower()
    return {
        ".amr": "audio/amr",
        ".wav": "audio/wav",
        ".mp3": "audio/mp3",
        ".opus": "audio/opus",
        ".ogg": "audio/ogg",
        ".flac": "audio/flac",
        ".m4a": "audio/mp4",
        ".aac": "audio/aac",
    }.get(ext, "audio/octet-stream")


# Extensions accepted by the agentscope OpenAIChatFormatter
_FORMATTER_SUPPORTED_AUDIO_EXTS = {".wav", ".mp3"}
_AMR_EXTENSIONS = (".amr", ".amr-wb")
_AMR_FFMPEG_PROBE_PARAMS = ["-analyzeduration", "200M", "-probesize", "200M"]
_WAV_AUDIO_CODEC = "pcm_s16le"


def _convert_audio_to_wav(src_path: str) -> Optional[str]:
    """Convert an audio file to .wav using ffmpeg if the extension is not
    natively supported by the LLM formatter.

    Uses a unique temporary file name to avoid overwriting existing files.

    Returns the path to the converted .wav file, or None if conversion
    failed or was not needed.
    """
    ext = (os.path.splitext(src_path)[1] or "").lower()
    if ext in _FORMATTER_SUPPORTED_AUDIO_EXTS:
        return None  # already supported, no conversion needed

    import subprocess
    import shutil
    import tempfile

    if not shutil.which("ffmpeg"):
        logger.warning(
            "ffmpeg not found; cannot convert %s audio to wav. "
            "Install ffmpeg to enable audio format conversion.",
            ext,
        )
        return None

    # Use a temp file in the same directory to avoid clobbering.
    src_dir = os.path.dirname(src_path) or "."
    fd, dst_path = tempfile.mkstemp(suffix=".wav", dir=src_dir)
    os.close(fd)

    # AMR (AMR-NB/AMR-WB) used by QQ voice messages has non-standard
    # encapsulation; increase analyzeduration and probesize so ffmpeg
    # can correctly detect the codec before decoding.
    amr_extra: list = (
        _AMR_FFMPEG_PROBE_PARAMS if ext in _AMR_EXTENSIONS else []
    )

    try:
        subprocess.run(
            [
                "ffmpeg",
                "-y",
                "-loglevel",
                "error",
                *amr_extra,
                "-i",
                src_path,
                "-acodec",
                _WAV_AUDIO_CODEC,
                "-ar",
                "16000",
                "-ac",
                "1",
                dst_path,
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            timeout=30,
            check=True,
        )
        logger.debug("Converted audio %s -> %s", src_path, dst_path)
        return dst_path
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
        stderr = getattr(e, "stderr", b"") or b""
        logger.warning(
            "Audio conversion failed for %s: %s\nffmpeg stderr: %s",
            src_path,
            e,
            stderr.decode(errors="replace"),
        )
        # Clean up the temp file on failure.
        try:
            os.unlink(dst_path)
        except OSError:
            pass
        return None


def _local_file_url(path: str) -> str:
    """Build a ``file://`` URL without percent-encoding non-ASCII chars."""
    return "file://" + str(Path(path).resolve())


def _update_block_with_local_path(
    block: dict,
    block_type: str,
    local_path: str,
) -> dict:
    """Update block with downloaded local path."""
    if block_type == "file":
        block["source"] = local_path
        if not block.get("filename"):
            block["filename"] = os.path.basename(local_path)
    else:
        if block_type == "audio":
            block["source"] = {
                "type": "url",
                "url": _local_file_url(local_path),
                "media_type": _media_type_from_path(local_path),
            }
        else:
            block["source"] = {
                "type": "url",
                "url": _local_file_url(local_path),
            }
    return block


def _handle_download_failure(block_type: str) -> Optional[dict]:
    """Handle download failure based on block type."""
    if block_type == "file":
        return {
            "type": "text",
            "text": "[Error: Unknown file source type or empty data]",
        }
    logger.debug("Failed to download %s block, keeping original", block_type)
    return None


async def _process_audio_block(
    message_content: list,
    index: int,
    local_path: str,
    block: dict,
) -> bool:
    """Handle an audio block according to the configured audio_mode.

    Modes:
      - ``"auto"`` (default): try transcription; if it succeeds, replace
        the audio block with the transcribed text and suppress file
        metadata.  If transcription fails (no provider, missing deps,
        API error), show a file-uploaded placeholder instead.  Audio is
        never sent directly to the model in this mode.
      - ``"native"``: send the audio block directly to the model
        (convert via ffmpeg if needed).  No transcription is attempted.
        If the file format is unsupported and conversion fails, a text
        placeholder is shown instead.

    Returns:
        True if the audio was fully handled (transcribed or sent natively)
        — the "file downloaded" notification will be suppressed.
        False if transcription failed — the notification is kept so the
        LLM knows the file path.
    """
    from .audio_transcription import transcribe_audio

    audio_mode = load_config().agents.audio_mode

    if audio_mode == "native":
        converted = await asyncio.to_thread(
            _convert_audio_to_wav,
            local_path,
        )
        ext = (os.path.splitext(local_path)[1] or "").lower()
        if converted:
            audio_path = converted
        elif ext in _FORMATTER_SUPPORTED_AUDIO_EXTS:
            # Already a supported format, no conversion needed.
            audio_path = local_path
        else:
            # Unsupported format and conversion failed — show placeholder
            # instead of sending an unsupported audio block to the model.
            message_content[index] = {
                "type": "text",
                "text": (
                    "[Voice message]: (audio conversion failed, "
                    "install ffmpeg to enable native audio)"
                ),
            }
            return True
        block["source"] = {
            "type": "url",
            "url": _local_file_url(audio_path),
            "media_type": _media_type_from_path(audio_path),
        }
        return True

    # "auto": attempt transcription.
    text = await transcribe_audio(local_path)
    if text:
        message_content[index] = {
            "type": "text",
            "text": f"[Voice message]: {text}",
        }
        return True

    # Transcription failed — show file-uploaded placeholder.
    message_content[index] = {
        "type": "text",
        "text": "[Voice message]: (audio file received)",
    }
    return False


async def _process_single_block(
    message_content: list,
    index: int,
    block: dict,
) -> Optional[str]:
    """
    Process a single file or media block.

    Returns:
        Optional[str]: The local path if download was successful,
        None otherwise.
    """
    block_type = block.get("type")
    if not isinstance(block_type, str):
        return None

    source, filename = _extract_source_and_filename(block, block_type)
    if source is None:
        return None

    # Normalize: when source is "base64" but data is a local path (e.g.
    # DingTalk voice returns path), treat as url only if under allowed dir.
    if (
        block_type == "audio"
        and isinstance(source, dict)
        and source.get("type") == "base64"
    ):
        data = source.get("data")
        if isinstance(data, str) and os.path.isfile(data):
            block["source"] = {
                "type": "url",
                "url": _local_file_url(data),
                "media_type": _media_type_from_path(data),
            }
            source = block["source"]

    try:
        local_path = await _process_single_file_block(source, filename)

        if local_path:
            if block_type == "audio":
                # Audio blocks need transcription or format conversion
                # depending on the configured audio_mode.
                _update_block_with_local_path(block, block_type, local_path)
                handled = await _process_audio_block(
                    message_content,
                    index,
                    local_path,
                    block,
                )
                if handled:
                    # Audio was transcribed or sent natively; suppress the
                    # "file downloaded" notification that would follow.
                    return None
            else:
                message_content[index] = _update_block_with_local_path(
                    block,
                    block_type,
                    local_path,
                )
            logger.debug(
                "Updated %s block with local path: %s",
                block_type,
                local_path,
            )
            return local_path
        else:
            error_block = _handle_download_failure(block_type)
            if error_block:
                message_content[index] = error_block
            return None

    except Exception as e:
        logger.error("Failed to process %s block: %s", block_type, e)
        if block_type == "file":
            message_content[index] = {
                "type": "text",
                "text": f"[Error: Failed to download file - {e}]",
            }
        return None


# pylint: disable=too-many-return-statements
def _coerce_block_to_dict(
    block,
) -> dict | None:
    """Convert a Pydantic block (or dict) to a dict for processing.

    Returns ``None`` for blocks that are not file/media types.
    For 2.0 ``DataBlock``, maps ``type="data"`` back to the concrete
    media category (``"image"``/``"audio"``/``"video"``) so downstream
    helpers recognise it.
    """
    if isinstance(block, dict):
        return block

    btype = getattr(block, "type", None)
    if btype == "data":
        source = getattr(block, "source", None)
        if source is None:
            return None
        mt = getattr(source, "media_type", "") or ""
        main = mt.split("/")[0]
        if main not in ("image", "audio", "video"):
            # Generic data block (e.g. application/*) — treat as file
            main = "file"
        src_dict: dict = {}
        src_type = getattr(source, "type", None)
        if src_type == "url":
            url_str = str(getattr(source, "url", ""))
            if url_str.startswith("file://"):
                url_str = url_str.removeprefix("file://")
            src_dict = {"type": "url", "url": url_str, "media_type": mt}
        elif src_type == "base64":
            src_dict = {
                "type": "base64",
                "data": getattr(source, "data", ""),
                "media_type": mt,
            }
        else:
            return None
        return {
            "type": main,
            "source": src_dict,
            "filename": getattr(block, "name", None),
        }

    if btype in ("file", "image", "audio", "video"):
        if hasattr(block, "model_dump"):
            return block.model_dump()
        return None

    return None


async def process_file_and_media_blocks_in_message(msg) -> None:
    """Process file and media blocks (file, image, audio, video) in messages.

    Downloads to local and updates paths/URLs.  Handles both dict blocks
    (1.x) and Pydantic block objects (2.0 ``DataBlock``).
    """
    messages = (
        [msg] if isinstance(msg, Msg) else msg if isinstance(msg, list) else []
    )

    for message in messages:
        if not isinstance(message, Msg):
            continue

        if not isinstance(message.content, list):
            continue

        downloaded_files = []

        for i, block in enumerate(message.content):
            # === 2.0 Pydantic DataBlock fast-path ===
            # Console uploads land as ``DataBlock(URLSource(url=file:///..))``
            # already pointing at ``media_dir``.  Skip the dict-based
            # download / mutation entirely (it would replace the Pydantic
            # block with a dict and break ``msg.has_content_blocks(...)``
            # in agentscope ``_handle_incoming_messages``).  We only need
            # the local path for the sibling-text-block injection below.
            if not isinstance(block, dict):
                source = getattr(block, "source", None)
                url = str(getattr(source, "url", "")) if source else ""
                if url.startswith("file://"):
                    local_path = url.removeprefix("file://")
                    downloaded_files.append((i, local_path))
                # Remote URL or no URL on a Pydantic block: skip silently.
                # Adding remote-download for Pydantic DataBlock is a
                # separate feature (would need to also convert the dict
                # result back into a DataBlock to preserve the array
                # type homogeneity that 2.0 expects).
                continue

            # === 1.x legacy dict path ===
            block_dict = _coerce_block_to_dict(block)
            if block_dict is None:
                continue

            block_type = block_dict.get("type")
            if block_type not in ["file", "image", "audio", "video"]:
                continue

            local_path = await _process_single_block(
                message.content,
                i,
                block_dict,
            )
            if local_path:
                downloaded_files.append((i, local_path))

        if downloaded_files:
            lang = load_config().agents.language
            for i, local_path in reversed(downloaded_files):
                text = (
                    f"用户上传文件，已经下载到 {local_path}"
                    if lang == "zh"
                    else f"User uploaded a file, downloaded to {local_path}"
                )
                message.content.insert(
                    i + 1,
                    TextBlock(type="text", text=text),
                )


def is_first_user_interaction(messages: list) -> bool:
    """Check if this is the first user interaction.

    Args:
        messages: List of Msg objects from memory.

    Returns:
        bool: True if this is the first user message with no assistant
              responses.
    """
    system_prompt_count = sum(1 for msg in messages if msg.role == "system")
    non_system_messages = messages[system_prompt_count:]

    user_msg_count = sum(
        1 for msg in non_system_messages if msg.role == "user"
    )
    assistant_msg_count = sum(
        1 for msg in non_system_messages if msg.role == "assistant"
    )

    return user_msg_count == 1 and assistant_msg_count == 0


def prepend_to_message_content(msg, guidance: str) -> None:
    """Prepend guidance text to message content.

    Handles both dict blocks and Pydantic TextBlock objects.
    """
    if isinstance(msg.content, str):
        msg.content = guidance + "\n\n" + msg.content
        return

    if not isinstance(msg.content, list):
        return

    for block in msg.content:
        btype = (
            block.get("type")
            if isinstance(block, dict)
            else getattr(block, "type", None)
        )
        if btype == "text":
            if isinstance(block, dict):
                block["text"] = guidance + "\n\n" + block.get("text", "")
            else:
                block.text = guidance + "\n\n" + (block.text or "")
            return

    msg.content.insert(0, TextBlock(type="text", text=guidance))
