"""zyx.processing.multimodal"""

from __future__ import annotations

from functools import lru_cache
from enum import Enum
from io import BytesIO
import mimetypes
import re
from urllib.parse import urlparse

import httpx
import urllib.request
import urllib.error
from pathlib import Path
from typing import Any, TypeAlias, Self, Tuple, TYPE_CHECKING

from pydantic_ai.messages import (
    AudioUrl,
    BinaryContent,
    DocumentUrl,
    ImageUrl,
    VideoUrl,
    UserContent,
)

from ._toon import object_as_text

if TYPE_CHECKING:
    from markitdown import MarkItDown

__all__ = (
    "MultimodalContent",
    "MultimodalContentOrigin",
    "MultimodalContentMediaType",
    "classify_multimodal_source",
    "render_multimodal_source_as_user_content",
    "render_multimodal_source_as_text",
    "render_multimodal_source_as_description",
)


MultimodalContent: TypeAlias = Path | str | bytes | Any
"""
The source location from where a piece of external content can originate from.

This can be one of the following types:
- A `Path` object representing a local file
- A string representing a URL to a text or multimodal file/html page
- A bytes object representing raw binary data
- Any non-string python object that can be encoded to text using the TOON format
"""


_URIS = ("http://", "https://", "ftp://")
_TEXT_EXTENSIONS: frozenset[str] = frozenset(
    (".txt", ".md", ".rst", ".json", ".yaml", ".yml", ".ini", ".toml", ".log")
)
_DOCUMENT_MIMETYPES: frozenset[str] = frozenset(
    (
        "application/pdf",
        "application/msword",
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "application/vnd.ms-powerpoint",
        "application/vnd.openxmlformats-officedocument.presentationml.presentation",
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    )
)
_URL_HEADERS = {
    "User-Agent": "zyx/1.0",
    "Accept": "*/*",
}


def _looks_like_url(value: str) -> bool:
    if not value or " " in value:
        return False
    if value.startswith(_URIS):
        return True
    # Avoid misclassifying local filesystem paths like "foo/bar.txt"
    if Path(value).exists():
        return False
    candidate = _normalize_url(value)
    try:
        parsed = httpx.URL(candidate)
    except Exception:
        return False
    host = parsed.host
    if not host:
        return False
    if host == "localhost":
        return True
    return "." in host
    return True


def _normalize_url(value: str) -> str:
    if value.startswith(_URIS):
        return value
    parsed = urlparse(value)
    if parsed.scheme:
        return value
    return f"https://{value}"


class MultimodalContentOrigin(Enum):
    """
    The computed origin of a snippet.
    """

    FILE = "file"
    URL = "url"
    STRING = "string"
    BYTES = "bytes"
    OBJECT = "object"

    @classmethod
    def classify(cls, source: MultimodalContent) -> Self:
        """
        Classifies the the origin of a snippet's source.
        """
        if isinstance(source, (bytes, bytearray, memoryview)):
            return cls.BYTES  # type: ignore

        elif isinstance(source, Path):
            if source.is_file():
                return cls.FILE  # type: ignore

        else:
            if isinstance(source, str):
                if source.startswith(_URIS):
                    return cls.URL  # type: ignore
                if _looks_like_url(source):
                    return cls.URL  # type: ignore
                if Path(source).exists() and Path(source).is_file():
                    return cls.FILE  # type: ignore

            else:
                return cls.OBJECT  # type: ignore

        return cls.STRING  # type: ignore


class MultimodalContentMediaType(Enum):
    """
    The media/content type of a source.
    """

    TEXT = "text/plain"
    HTML = "text/html"
    DOCUMENT = "document"
    IMAGE = "image"
    AUDIO = "audio"
    VIDEO = "video"
    UNKNOWN = "unknown"

    @classmethod
    def classify(
        cls, source: MultimodalContent, origin: MultimodalContentOrigin
    ) -> Self:
        """
        Classifies the the type of a snippet's content.
        """
        if (
            origin == MultimodalContentOrigin.OBJECT
            or origin == MultimodalContentOrigin.STRING
        ):
            # objects are encoded to text using the TOON format
            return cls.TEXT  # type: ignore

        if origin == MultimodalContentOrigin.BYTES:
            data = (
                source.tobytes()
                if isinstance(source, memoryview)
                else bytes(source)  # type: ignore
            )
            return _classify_magic_bytes(data)  # type: ignore

        path_text = ""
        suffix = ""
        if origin == MultimodalContentOrigin.FILE:
            path = source if isinstance(source, Path) else Path(source)  # type: ignore
            path_text = str(path)
            suffix = path.suffix.lower()
        elif origin == MultimodalContentOrigin.URL and isinstance(source, str):
            path_text = re.split(r"[?#]", source, 1)[0]
            suffix = Path(path_text).suffix.lower()

        if suffix in _TEXT_EXTENSIONS:
            return cls.TEXT  # type: ignore
        if suffix in (".html", ".htm", ".xhtml"):
            return cls.HTML  # type: ignore
        if suffix in (
            ".pdf",
            ".doc",
            ".docx",
            ".ppt",
            ".pptx",
            ".xls",
            ".xlsx",
        ):
            return cls.DOCUMENT  # type: ignore

        mime_type, _ = mimetypes.guess_type(path_text)
        if mime_type == "text/html":
            return cls.HTML  # type: ignore
        if mime_type in _DOCUMENT_MIMETYPES:
            return cls.DOCUMENT  # type: ignore
        if mime_type:
            if mime_type.startswith("image/"):
                return cls.IMAGE  # type: ignore
            if mime_type.startswith("audio/"):
                return cls.AUDIO  # type: ignore
            if mime_type.startswith("video/"):
                return cls.VIDEO  # type: ignore
            if mime_type.startswith("text/"):
                return cls.TEXT  # type: ignore

        if (
            origin == MultimodalContentOrigin.URL
            and not suffix
            and not mime_type
        ):
            detected, _ = _detect_url_media_type(source)
            if detected:
                return detected  # type: ignore
            # Default unknown URLs to text to avoid treating common web pages as UNKNOWN.
            return cls.TEXT  # type: ignore

        return cls.UNKNOWN  # type: ignore


def _classify_content_type(
    content_type: str | None,
) -> MultimodalContentMediaType | None:
    if not content_type:
        return None
    ct = content_type.split(";", 1)[0].strip().lower()
    if not ct:
        return None
    if ct == "text/html":
        return MultimodalContentMediaType.HTML
    if ct in _DOCUMENT_MIMETYPES or ct == "application/pdf":
        return MultimodalContentMediaType.DOCUMENT
    if ct.startswith("image/"):
        return MultimodalContentMediaType.IMAGE
    if ct.startswith("audio/"):
        return MultimodalContentMediaType.AUDIO
    if ct.startswith("video/"):
        return MultimodalContentMediaType.VIDEO
    if ct.startswith("text/"):
        return MultimodalContentMediaType.TEXT
    return None


def _classify_magic_bytes(
    data: bytes,
) -> MultimodalContentMediaType | None:
    head = data[:512]
    if head.startswith(b"%PDF"):
        return MultimodalContentMediaType.DOCUMENT
    if head.startswith(b"\x89PNG\r\n\x1a\n") or head[:3] == b"\xff\xd8\xff":
        return MultimodalContentMediaType.IMAGE
    if head.startswith(b"GIF87a") or head.startswith(b"GIF89a"):
        return MultimodalContentMediaType.IMAGE
    if head[:4] == b"RIFF" and head[8:12] == b"WEBP":
        return MultimodalContentMediaType.IMAGE
    if head[:4] == b"RIFF" and head[8:12] == b"WAVE":
        return MultimodalContentMediaType.AUDIO
    if (
        head.startswith(b"OggS")
        or head.startswith(b"fLaC")
        or head.startswith(b"ID3")
    ):
        return MultimodalContentMediaType.AUDIO
    if len(head) >= 8 and head[4:8] == b"ftyp":
        return MultimodalContentMediaType.VIDEO
    if head.startswith(b"\x1a\x45\xdf\xa3"):
        return MultimodalContentMediaType.VIDEO
    if b"\x00" in head:
        return MultimodalContentMediaType.UNKNOWN
    try:
        text = head.decode("utf-8")
    except UnicodeDecodeError:
        return MultimodalContentMediaType.UNKNOWN
    if re.search(r"(?i)<!doctype html|<html\\b|<head\\b|<body\\b", text):
        return MultimodalContentMediaType.HTML
    return MultimodalContentMediaType.TEXT


@lru_cache(maxsize=128)
def _detect_url_media_type(
    url: str,
) -> Tuple[MultimodalContentMediaType | None, str | None]:
    url = _normalize_url(url)
    try:
        req = urllib.request.Request(url, method="HEAD", headers=_URL_HEADERS)
        with urllib.request.urlopen(req, timeout=5) as resp:
            content_type = resp.headers.get("Content-Type")
            classified = _classify_content_type(content_type)
            return classified, (
                content_type.split(";", 1)[0].strip().lower()
                if content_type
                else None
            )
    except (urllib.error.URLError, urllib.error.HTTPError):
        pass

    try:
        req = urllib.request.Request(
            url,
            method="GET",
            headers={**_URL_HEADERS, "Range": "bytes=0-511"},
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            content_type = resp.headers.get("Content-Type")
            ct_class = _classify_content_type(content_type)
            if ct_class:
                return ct_class, (
                    content_type.split(";", 1)[0].strip().lower()
                    if content_type
                    else None
                )
            data = resp.read(512)
            return _classify_magic_bytes(data), None
    except (urllib.error.URLError, urllib.error.HTTPError):
        return None, None


def classify_multimodal_source(
    source: MultimodalContent,
) -> Tuple[MultimodalContentOrigin, MultimodalContentMediaType]:
    """
    Classifies a multimodal source into its origin and type.
    """
    origin = MultimodalContentOrigin.classify(source)
    return origin, MultimodalContentMediaType.classify(source, origin)


@lru_cache(maxsize=1)
def _get_markdown_converter() -> MarkItDown:
    from markitdown import MarkItDown

    return MarkItDown()


def _render_multimodal_file_or_url_as_text(source: str) -> str:
    try:
        return _get_markdown_converter().convert(source).markdown
    except Exception as e:
        raise ValueError(f"Failed to render file or URL as text: {e}") from e


def render_multimodal_source_as_user_content(
    source: MultimodalContent,
    origin: MultimodalContentOrigin,
    media_type: MultimodalContentMediaType,
) -> UserContent:
    """
    Renders a multimodal source as a compatible `pydantic-ai` UserContent object.
    """
    if origin == MultimodalContentOrigin.OBJECT:
        return object_as_text(source)

    if origin == MultimodalContentOrigin.STRING:
        if media_type == MultimodalContentMediaType.HTML:
            from markitdown import StreamInfo

            buffer = BytesIO(str(source).encode("utf-8"))
            info = StreamInfo(mimetype="text/html", extension=".html")
            return (
                _get_markdown_converter()
                .convert(buffer, stream_info=info)
                .markdown
            )
        return str(source)

    if origin == MultimodalContentOrigin.URL and isinstance(source, str):
        normalized = _normalize_url(source)
        guessed_media_type, _ = mimetypes.guess_type(normalized)
        if not guessed_media_type:
            _, detected = _detect_url_media_type(normalized)
            guessed_media_type = detected
        if media_type == MultimodalContentMediaType.IMAGE:
            return ImageUrl(url=normalized, media_type=guessed_media_type)
        if media_type == MultimodalContentMediaType.AUDIO:
            return AudioUrl(url=normalized, media_type=guessed_media_type)
        if media_type == MultimodalContentMediaType.VIDEO:
            return VideoUrl(url=normalized, media_type=guessed_media_type)
        if media_type in (
            MultimodalContentMediaType.TEXT,
            MultimodalContentMediaType.HTML,
            MultimodalContentMediaType.DOCUMENT,
        ):
            return _render_multimodal_file_or_url_as_text(normalized)
        return (
            DocumentUrl(url=normalized, media_type=guessed_media_type)
            if guessed_media_type
            else normalized
        )

    if origin == MultimodalContentOrigin.FILE:
        path = source if isinstance(source, Path) else Path(source)  # type: ignore
        if media_type in (
            MultimodalContentMediaType.TEXT,
            MultimodalContentMediaType.HTML,
            MultimodalContentMediaType.DOCUMENT,
        ):
            return _render_multimodal_file_or_url_as_text(str(path))
        guessed_media_type, _ = mimetypes.guess_type(str(path))
        data = path.read_bytes()
        return BinaryContent.narrow_type(
            BinaryContent(
                data=data,
                media_type=guessed_media_type or "application/octet-stream",
            )
        )

    if origin == MultimodalContentOrigin.BYTES:
        data = (
            source.tobytes()
            if isinstance(source, memoryview)
            else bytes(source)  # type: ignore
        )
        if media_type == MultimodalContentMediaType.HTML:
            from markitdown import StreamInfo

            buffer = BytesIO(data)
            info = StreamInfo(mimetype="text/html", extension=".html")
            return (
                _get_markdown_converter()
                .convert(buffer, stream_info=info)
                .markdown
            )
        if media_type == MultimodalContentMediaType.TEXT:
            return data.decode("utf-8")
        return BinaryContent.narrow_type(
            BinaryContent(data=data, media_type="application/octet-stream")
        )

    return str(source)


def render_multimodal_source_as_text(
    source: MultimodalContent,
    origin: MultimodalContentOrigin,
    media_type: MultimodalContentMediaType,
) -> str:
    """
    Renders a multimodal source as text, or a textual description for
    multimodal content.
    """
    if origin == MultimodalContentOrigin.OBJECT:
        return object_as_text(source)

    if media_type in (
        MultimodalContentMediaType.TEXT,
        MultimodalContentMediaType.HTML,
        MultimodalContentMediaType.DOCUMENT,
        MultimodalContentMediaType.UNKNOWN,
    ):
        if origin == MultimodalContentOrigin.STRING and isinstance(
            source, str
        ):
            return source

        if origin == MultimodalContentOrigin.FILE:
            path = source if isinstance(source, Path) else Path(source)  # type: ignore
            return _render_multimodal_file_or_url_as_text(str(path))

        if origin == MultimodalContentOrigin.URL and isinstance(source, str):
            return _render_multimodal_file_or_url_as_text(
                _normalize_url(source)
            )

        if origin == MultimodalContentOrigin.BYTES:
            data = (
                source.tobytes()
                if isinstance(source, memoryview)
                else bytes(source)  # type: ignore
            )
            if media_type == MultimodalContentMediaType.TEXT:
                return data.decode("utf-8")
            if media_type == MultimodalContentMediaType.HTML:
                from markitdown import StreamInfo

                buffer = BytesIO(data)
                info = StreamInfo(mimetype="text/html", extension=".html")
                return (
                    _get_markdown_converter()
                    .convert(buffer, stream_info=info)
                    .markdown
                )
            return "[document bytes]"

    if origin == MultimodalContentOrigin.URL and isinstance(source, str):
        return f"[{media_type.value} content at {_normalize_url(source)}]"
    if origin == MultimodalContentOrigin.FILE:
        path = source if isinstance(source, Path) else Path(source)  # type: ignore
        return f"[{media_type.value} content from {path}]"
    if origin == MultimodalContentOrigin.BYTES:
        return f"[{media_type.value} content bytes]"

    return str(source)


def render_multimodal_source_as_description(
    source: MultimodalContent,
    origin: MultimodalContentOrigin,
    media_type: MultimodalContentMediaType,
) -> str:
    """
    Renders a multimodal source as a description, or a textual description for
    multimodal content.
    """
    if origin == MultimodalContentOrigin.OBJECT:
        return f"[object of type {type(source).__name__}]"

    if origin == MultimodalContentOrigin.STRING:
        length = len(source) if isinstance(source, str) else 0
        return f"[string value, length={length}]"

    if origin == MultimodalContentOrigin.URL:
        return f"[URL: {source} | type={media_type.value}]"

    if origin == MultimodalContentOrigin.FILE:
        path = source if isinstance(source, Path) else Path(source)  # type: ignore[arg-type]
        name = path.name or "file"
        return f"[{media_type.value} file: {name}]"

    if origin == MultimodalContentOrigin.BYTES:
        size = (
            len(source)
            if isinstance(source, (bytes, bytearray))
            else len(bytes(source))  # type: ignore
        )
        return f"[{media_type.value} bytes, size={size}]"

    return f"[{media_type.value} content]"
