from __future__ import annotations

from collections.abc import Iterable
from pathlib import Path
from typing import Any

IMAGE_SUFFIXES = {".avif", ".bmp", ".gif", ".jpeg", ".jpg", ".png", ".webp"}
TEXT_SUFFIXES = {
    ".bat",
    ".c",
    ".cfg",
    ".cpp",
    ".css",
    ".csv",
    ".go",
    ".h",
    ".hpp",
    ".html",
    ".ini",
    ".java",
    ".js",
    ".json",
    ".jsonl",
    ".jsx",
    ".kt",
    ".log",
    ".md",
    ".php",
    ".ps1",
    ".py",
    ".rb",
    ".rs",
    ".rst",
    ".sh",
    ".sql",
    ".swift",
    ".tex",
    ".toml",
    ".ts",
    ".tsv",
    ".tsx",
    ".txt",
    ".xml",
    ".yaml",
    ".yml",
}
MAX_ATTACHMENTS = 2
MAX_TEXT_FILE_BYTES = 64 * 1024


def file_path(value: Any) -> str | None:
    """Return a file path from a Gradio file value."""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        path = value.get("path") or value.get("name")
        if path:
            return str(path)
        nested_file = value.get("file")
        return file_path(nested_file) if nested_file is not None else None
    path = getattr(value, "path", None) or getattr(value, "name", None)
    return str(path) if path else None


def _text_file_part(path: str) -> dict[str, str]:
    file = Path(path)
    try:
        if file.stat().st_size > MAX_TEXT_FILE_BYTES:
            raise ValueError(
                f"{file.name} is too large. Text files must be at most "
                f"{MAX_TEXT_FILE_BYTES // 1024} KB."
            )
        raw = file.read_bytes()
    except OSError as exc:
        raise ValueError(f"Could not read attached text file: {file.name}") from exc
    if len(raw) > MAX_TEXT_FILE_BYTES:
        raise ValueError(
            f"{file.name} is too large. Text files must be at most "
            f"{MAX_TEXT_FILE_BYTES // 1024} KB."
        )
    if b"\x00" in raw:
        raise ValueError(f"{file.name} is not a plain-text file.")
    try:
        text = raw.decode("utf-8-sig").replace("\r\n", "\n").replace("\r", "\n")
    except UnicodeDecodeError as exc:
        raise ValueError(f"{file.name} must use UTF-8 text encoding.") from exc
    body = text.strip() or "(empty file)"
    return {
        "type": "text",
        "text": (
            f"--- BEGIN ATTACHED TEXT FILE: {file.name} ---\n"
            f"{body}\n"
            f"--- END ATTACHED TEXT FILE: {file.name} ---"
        ),
    }


def attachment_parts(
    message: dict[str, Any],
    limit: int = MAX_ATTACHMENTS,
) -> list[dict[str, str]]:
    """Validate image or UTF-8 text attachments and return model content parts."""
    raw_files = message.get("files") or []
    paths = [path for item in raw_files if (path := file_path(item))]
    if len(paths) > limit:
        raise ValueError(f"Attach at most {limit} image or text files per message.")

    parts: list[dict[str, str]] = []
    for path in paths:
        suffix = Path(path).suffix.lower()
        if suffix in IMAGE_SUFFIXES:
            parts.append({"type": "image", "path": path})
        elif suffix in TEXT_SUFFIXES:
            parts.append(_text_file_part(path))
        else:
            raise ValueError("Only image and plain-text files are supported.")
    return parts


def _content_parts(content: Any) -> list[dict[str, str]]:
    if isinstance(content, str) and content.strip():
        return [{"type": "text", "text": content.strip()}]
    if isinstance(content, dict):
        path = file_path(content)
        if path:
            suffix = Path(path).suffix.lower()
            if suffix in IMAGE_SUFFIXES:
                return [{"type": "image", "path": path}]
            if suffix in TEXT_SUFFIXES:
                return [_text_file_part(path)]
        value = content.get("text") or content.get("value")
        if isinstance(value, str) and value.strip():
            return [{"type": "text", "text": value.strip()}]
    return []


def build_messages(
    message: dict[str, Any],
    history: Iterable[dict[str, Any]],
    system_prompt: str,
    history_limit: int,
) -> list[dict[str, Any]]:
    """Convert Gradio chat data to the Transformers multimodal chat format."""
    messages: list[dict[str, Any]] = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})

    valid_history = [
        item
        for item in history
        if isinstance(item, dict) and item.get("role") in {"user", "assistant"}
    ][-history_limit:]
    for item in valid_history:
        parts = _content_parts(item.get("content"))
        if parts:
            messages.append({"role": item["role"], "content": parts})

    current_parts = attachment_parts(message)
    has_image = any(part["type"] == "image" for part in current_parts)
    has_text_file = any(part["type"] == "text" for part in current_parts)
    text = str(message.get("text") or "").strip()
    if text:
        current_parts.append({"type": "text", "text": text})
    elif has_image and has_text_file:
        current_parts.append({"type": "text", "text": "Analyze the attached files."})
    elif has_image:
        current_parts.append({"type": "text", "text": "Describe this image in detail."})
    elif has_text_file:
        current_parts.append({"type": "text", "text": "Summarize the attached text."})
    else:
        raise ValueError("Enter a message or attach an image or text file.")

    messages.append({"role": "user", "content": current_parts})
    return messages


def estimate_duration(
    _message: Any,
    _history: Any,
    _system_prompt: str,
    _thinking: bool,
    max_new_tokens: int,
    *_args: Any,
    **_kwargs: Any,
) -> int:
    """Estimate a bounded ZeroGPU allocation from the requested output length."""
    return min(600, 55 + round(int(max_new_tokens) * 0.32))