# -*- coding: utf-8 -*-
import os
import base64
import hashlib
import logging
import shutil
import subprocess
import urllib.parse
from typing import Optional, Any
from pathlib import Path

logger = logging.getLogger(__name__)

# Global token counter instance (lazy initialization)
_token_counter = None


async def download_file_from_base64(
    base64_data: str,
    filename: Optional[str] = None,
    download_dir: str = "downloads",
) -> str:
    """
    Save base64-encoded file data to local download directory.

    Args:
        base64_data (`str`):
            Base64-encoded file content.
        filename (`str`, optional):
            The filename to save. If not provided, will generate one.
        download_dir (`str`):
            The directory to save files. Defaults to "downloads".

    Returns:
        `str`:
            The local file path.
    """
    try:
        # Decode base64 data
        file_content = base64.b64decode(base64_data)

        # Create download directory if not exists
        download_path = Path(download_dir)
        download_path.mkdir(parents=True, exist_ok=True)

        # Generate filename if not provided
        if not filename:
            # Use hash of content as filename
            file_hash = hashlib.md5(file_content).hexdigest()
            filename = f"file_{file_hash}"

        # Save file
        local_file_path = download_path / filename
        with open(local_file_path, "wb") as f:
            f.write(file_content)

        logger.info("Downloaded file to: %s", local_file_path)
        return str(local_file_path.absolute())

    except Exception as e:
        logger.error("Failed to download file from base64: %s", e)
        raise


async def download_file_from_url(
    url: str,
    filename: Optional[str] = None,
    download_dir: str = "downloads",
) -> str:
    """
    Download a file from URL to local download directory using wget or curl.

    Args:
        url (`str`):
            The URL of the file to download.
        filename (`str`, optional):
            The filename to save. If not provided, will extract from URL or
            generate a hash-based name.
        download_dir (`str`):
            The directory to save files. Defaults to "downloads".

    Returns:
        `str`:
            The local file path.
    """
    try:
        # Create download directory if not exists
        download_path = Path(download_dir)
        download_path.mkdir(parents=True, exist_ok=True)

        # Generate filename if not provided
        if not filename:
            # Try to extract filename from URL
            parsed_url = urllib.parse.urlparse(url)
            url_filename = os.path.basename(parsed_url.path)
            if url_filename:
                filename = url_filename
            else:
                # Use hash of URL as filename
                url_hash = hashlib.md5(url.encode()).hexdigest()
                filename = f"file_{url_hash}"

        # Full local file path
        local_file_path = download_path / filename

        # Try wget first, then fall back to curl
        download_success = False
        try:
            subprocess.run(
                ["wget", "-q", "-O", str(local_file_path), url],
                capture_output=True,
                timeout=60,
                check=True,
            )
            logger.info("Downloaded file via wget to: %s", local_file_path)
            download_success = True
        except (subprocess.CalledProcessError, FileNotFoundError) as e:
            # wget failed or not available, try curl
            logger.debug("wget failed, trying curl: %s", e)
            try:
                subprocess.run(
                    ["curl", "-s", "-L", "-o", str(local_file_path), url],
                    capture_output=True,
                    timeout=60,
                    check=True,
                )
                logger.info("Downloaded file via curl to: %s", local_file_path)
                download_success = True
            except (
                subprocess.CalledProcessError,
                FileNotFoundError,
            ) as curl_err:
                logger.error(
                    "Both wget and curl failed for URL %s: %s",
                    url,
                    curl_err,
                )
                raise RuntimeError(
                    "Failed to download file: both wget and curl failed",
                ) from curl_err

        # Verify file was downloaded
        if not download_success or not local_file_path.exists():
            raise FileNotFoundError("Downloaded file does not exist")

        if local_file_path.stat().st_size == 0:
            raise ValueError("Downloaded file is empty")

        return str(local_file_path.absolute())

    except subprocess.TimeoutExpired as e:
        logger.error("Download timeout for URL: %s", url)
        raise TimeoutError(f"Download timeout for URL: {url}") from e
    except Exception as e:
        logger.error("Failed to download file from URL %s: %s", url, e)
        raise


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

    Args:
        source (`dict`):
            The source dict containing file information.
        filename (`str`, optional):
            The filename to save.

    Returns:
        `str` or `None`:
            The local file path if successful, None otherwise.
    """
    # Handle Base64Source
    if isinstance(source, dict) and source.get("type") == "base64":
        # Check if 'data' key exists (allow empty string)
        if "data" in source:
            base64_data = source.get("data", "")
            local_path = await download_file_from_base64(
                base64_data,
                filename,
            )
            logger.info(
                "Processed base64 file block: %s -> %s",
                filename or "unnamed",
                local_path,
            )
            return local_path

    # Handle URLSource
    elif isinstance(source, dict) and source.get("type") == "url":
        # Check if 'url' key exists and is not empty
        url = source.get("url", "")
        if url:
            local_path = await download_file_from_url(
                url,
                filename,
            )
            logger.info(
                "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")

    # Media blocks (image, audio, video)
    source = block.get("source", {})
    if not isinstance(source, dict):
        return None, None

    # Try to extract filename from URL
    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 _update_block_with_local_path(
    block: dict,
    block_type: str,
    local_path: str,
) -> dict:
    """Update block with downloaded local path."""
    if block_type == "file":
        # File block: update source directly
        block["source"] = local_path
        if not block.get("filename"):
            block["filename"] = os.path.basename(local_path)
    else:
        # Media blocks: update source with URL structure
        block["source"] = {"type": "url", "url": Path(local_path).as_uri()}
    return block


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


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

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

    try:
        # Download to local
        local_path = await _process_single_file_block(source, filename)

        if local_path:
            # Update block with local path
            message_content[index] = _update_block_with_local_path(
                block,
                block_type,
                local_path,
            )
            logger.info(
                "Updated %s block with local path: %s",
                block_type,
                local_path,
            )
            return local_path
        else:
            # Handle download failure
            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)
        # File block: replace with error; Media blocks: keep original
        if block_type == "file":
            message_content[index] = {
                "type": "text",
                "text": f"[Error: Failed to download file - {e}]",
            }
        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.

    Args:
        msg: The message object (Msg or list[Msg]) to process.
    """
    from agentscope.message import Msg

    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

        # Collect download results with their indices
        downloaded_files = []

        # Process each content block
        for i, block in enumerate(message.content):
            if not isinstance(block, dict):
                continue

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

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

        # Add text blocks for successfully downloaded files
        for i, local_path in reversed(downloaded_files):
            text_block = {
                "type": "text",
                "text": f"用户上传文件，已经下载到 {local_path}",
            }
            # Insert the text block right after the file block
            message.content.insert(i + 1, text_block)


def copy_md_files(language: str) -> int:
    """Copy md files from agents/md_files to working directory.

    Args:
        language: Language code (e.g. 'en', 'zh')

    Returns:
        Number of md files copied
    """
    from ..constant import WORKING_DIR

    # Get md_files directory path with language subdirectory
    md_files_dir = Path(__file__).parent / "md_files" / language

    if not md_files_dir.exists():
        logger.warning(
            "MD files directory not found: %s, falling back to 'en'",
            md_files_dir,
        )
        # Fallback to English if specified language not found
        md_files_dir = Path(__file__).parent / "md_files" / "en"
        if not md_files_dir.exists():
            logger.error("Default 'en' md files not found either")
            return 0

    # Ensure working directory exists
    WORKING_DIR.mkdir(parents=True, exist_ok=True)

    # Copy all .md files to working directory
    copied_count = 0
    for md_file in md_files_dir.glob("*.md"):
        target_file = WORKING_DIR / md_file.name
        try:
            shutil.copy2(md_file, target_file)
            logger.info("Copied md file: %s", md_file.name)
            copied_count += 1
        except Exception as e:
            logger.error(
                "Failed to copy md file '%s': %s",
                md_file.name,
                e,
            )

    if copied_count > 0:
        logger.info(
            "Copied %d md file(s) [%s] to %s",
            copied_count,
            language,
            WORKING_DIR,
        )

    return copied_count


def _get_token_counter():
    """Get or initialize the global token counter instance.

    Returns:
        TokenCounterBase: The token counter instance for Qwen models.

    Raises:
        RuntimeError: If token counter initialization fails.
    """
    global _token_counter
    if _token_counter is None:
        from agentscope.token import HuggingFaceTokenCounter

        # Use Qwen tokenizer for DashScope models
        # Qwen3 series uses the same tokenizer as Qwen2.5
        _token_counter = HuggingFaceTokenCounter(
            pretrained_model_name_or_path="Qwen/Qwen2.5-7B-Instruct",
            use_mirror=True,  # Use HF mirror for users in China
            use_fast=True,
            trust_remote_code=True,
        )
        logger.info("Token counter initialized with Qwen tokenizer")
    return _token_counter


async def count_message_tokens(
    messages: list[dict],
    tools: list[dict] | None = None,
    **kwargs: Any,
) -> int:
    """Count tokens in messages using the token counter.

    This function uses the HuggingFace token counter for Qwen models.

    Args:
        messages: List of message dictionaries in chat format
        tools: Optional list of tool schemas for token counting
        **kwargs: Additional arguments passed to the tokenizer

    Returns:
        int: The number of tokens in the messages

    Raises:
        RuntimeError: If token counter fails to initialize or count.
    """
    token_counter = _get_token_counter()
    token_count = await token_counter.count(
        messages=messages,
        tools=tools,
        **kwargs,
    )
    logger.debug(
        "Counted %d tokens in %d messages",
        token_count,
        len(messages),
    )
    return token_count


def check_valid_messages(messages: list) -> bool:
    """
    Check if the messages are valid by ensuring all tool_use blocks have
    corresponding tool_result blocks.

    Args:
        messages: List of Msg objects to validate.

    Returns:
        bool: True if all tool_use IDs have matching tool_result IDs,
              False otherwise.
    """

    def _get_tool_ids(msgs: list, block_type: str) -> set[str]:
        """Get all block ids of specified type from messages."""
        ids = set()
        for msg in msgs:
            if not isinstance(msg.content, list):
                continue
            for block in msg.content:
                if isinstance(block, dict) and block.get("type") == block_type:
                    block_id = block.get("id")
                    if block_id:
                        ids.add(block_id)
        return ids

    tool_use_ids = _get_tool_ids(messages, "tool_use")
    tool_result_ids = _get_tool_ids(messages, "tool_result")
    return tool_use_ids == tool_result_ids


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.

    Args:
        msg: Msg object to modify.
        guidance: Text to prepend to the message content.
    """
    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:
        if isinstance(block, dict) and block.get("type") == "text":
            block["text"] = guidance + "\n\n" + block.get("text", "")
            return

    msg.content.insert(0, {"type": "text", "text": guidance})
