"""Image utilities for ``:::image`` block handling."""

from __future__ import annotations

import base64
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
from urllib.error import URLError
from urllib.parse import urlparse
from urllib.request import Request, urlopen

from dotpromptz.typing import MediaContent, MediaPart, Message, Part

_IMAGE_EXTENSION_TO_MIME: dict[str, str] = {
    '.gif': 'image/gif',
    '.jpeg': 'image/jpeg',
    '.jpg': 'image/jpeg',
    '.png': 'image/png',
    '.webp': 'image/webp',
}
_REMOTE_SCHEMES = frozenset({'http', 'https'})
_DATA_URI_PREFIX = 'data:'
_DEFAULT_REMOTE_TIMEOUT_SECONDS = 10.0


@dataclass(frozen=True)
class ResolvedImageRef:
    """Normalized image reference used by adapters and compiler.

    Attributes:
        kind: Source kind, either remote URL or local file.
        mime_type: Effective MIME type for this image.
        url: Remote URL when ``kind == 'remote_url'``.
        file_path: Local path when ``kind == 'local_file'``.
        data: Optional file bytes when ``read_bytes=True``.
    """

    kind: Literal['remote_url', 'local_file']
    mime_type: str
    url: str | None = None
    file_path: Path | None = None
    data: bytes | None = None


def resolve_image_value(
    lines: list[str],
    image_block_index: int,
    *,
    is_block_line: Callable[[str], bool],
) -> tuple[str, int] | None:
    """Resolve an ``:::image`` value from subsequent lines.

    The first non-empty line after the directive is used as image value.
    If another block line is encountered first, or EOF is reached, returns
    ``None``.

    Args:
        lines: Full source lines.
        image_block_index: Zero-based index of the ``:::image`` line.
        is_block_line: Predicate used to detect other ``:::`` directives.

    Returns:
        A tuple ``(value, line_index)`` where *line_index* is zero-based,
        or ``None`` when no value can be resolved.
    """
    for idx in range(image_block_index + 1, len(lines)):
        line = lines[idx]
        if not line.strip():
            continue
        if is_block_line(line):
            return None
        return line.strip(), idx
    return None


def infer_image_mime(value: str) -> tuple[str, bool]:
    """Infer image MIME type from path/URL file extension.

    Returns ``('image/png', True)`` when inference fails.

    Args:
        value: Rendered path or URL.

    Returns:
        A tuple ``(mime_type, used_default)``.
    """
    parsed = urlparse(value)
    candidate = parsed.path or value
    suffix = Path(candidate).suffix.lower()
    mime_type = _IMAGE_EXTENSION_TO_MIME.get(suffix)
    if mime_type is not None:
        return mime_type, False
    return 'image/png', True


def _prompt_base_dir(prompt_file_path: Path | None) -> Path:
    """Resolve the base directory used for relative image paths."""
    if prompt_file_path is None:
        return Path.cwd()
    if prompt_file_path.is_dir():
        return prompt_file_path
    return prompt_file_path.parent


def _looks_like_windows_path(value: str) -> bool:
    """Return whether ``value`` resembles a Windows absolute path."""
    return len(value) >= 3 and value[1] == ':' and value[0].isalpha() and value[2] in {'\\', '/'}


def _resolve_local_image_path(value: str, *, prompt_file_path: Path | None = None) -> Path:
    """Resolve a local image path (relative or absolute) to an absolute path."""
    raw_path = Path(value).expanduser()
    if raw_path.is_absolute():
        path = raw_path.resolve()
    else:
        base_dir = _prompt_base_dir(prompt_file_path)
        path = (base_dir / raw_path).resolve()
    return path


def _read_remote_image_bytes(url: str, *, timeout_seconds: float) -> bytes:
    """Read image bytes from a remote HTTP(S) URL."""
    request = Request(url, headers={'User-Agent': 'dotpromptz/1.0'})
    try:
        with urlopen(request, timeout=timeout_seconds) as response:
            data = response.read()
    except (OSError, URLError) as exc:
        raise ValueError(f'Failed to fetch image URL {url}: {exc}') from exc
    if not data:
        raise ValueError(f'Image URL returned empty content: {url}')
    return data


def resolve_image_ref(
    value: str,
    *,
    content_type: str | None = None,
    prompt_file_path: Path | None = None,
    read_bytes: bool = False,
    read_remote_bytes: bool = False,
    remote_timeout_seconds: float = _DEFAULT_REMOTE_TIMEOUT_SECONDS,
) -> ResolvedImageRef:
    """Resolve an image reference into remote URL or local file metadata.

    Args:
        value: Raw image reference from ``:::image`` or direct ``MediaPart``.
        content_type: Optional MIME type from parsed prompt metadata.
        prompt_file_path: Optional prompt path used for resolving relative paths.
        read_bytes: Whether local file bytes should be read and returned.
        read_remote_bytes: Whether remote URL bytes should also be downloaded.
        remote_timeout_seconds: Timeout used when downloading remote URL bytes.

    Returns:
        A normalized ``ResolvedImageRef``.

    Raises:
        ValueError: If reference is invalid, unsupported, or local file cannot be read.
    """
    stripped = value.strip()
    if not stripped:
        raise ValueError('Image reference cannot be empty.')
    if stripped.startswith(_DATA_URI_PREFIX):
        raise ValueError(
            'data: image URLs are not supported. '
            'Use a local file path (relative/absolute) or an http/https URL.',
        )

    parsed = urlparse(stripped)
    if parsed.scheme in _REMOTE_SCHEMES:
        mime_type = content_type or infer_image_mime(stripped)[0]
        data: bytes | None = None
        if read_bytes and read_remote_bytes:
            data = _read_remote_image_bytes(stripped, timeout_seconds=remote_timeout_seconds)
        return ResolvedImageRef(kind='remote_url', url=stripped, mime_type=mime_type, data=data)

    if parsed.scheme and not _looks_like_windows_path(stripped):
        raise ValueError(
            f'Unsupported image reference scheme: {parsed.scheme!r}. '
            'Use a local file path or an http/https URL.',
        )

    path = _resolve_local_image_path(stripped, prompt_file_path=prompt_file_path)
    if not path.exists():
        raise ValueError(f'Image file does not exist: {path}')
    if not path.is_file():
        raise ValueError(f'Image path must be a file: {path}')

    mime_type = content_type or infer_image_mime(path.as_posix())[0]
    data: bytes | None = None
    if read_bytes:
        try:
            data = path.read_bytes()
        except OSError as exc:
            raise ValueError(f'Failed to read image file {path}: {exc}') from exc
    return ResolvedImageRef(kind='local_file', file_path=path, mime_type=mime_type, data=data)


def to_data_url(data: bytes, mime_type: str) -> str:
    """Encode image bytes into a data URL string."""
    encoded = base64.b64encode(data).decode('ascii')
    return f'data:{mime_type};base64,{encoded}'


def normalize_media_messages(
    messages: list[Message],
    *,
    prompt_file_path: Path | None = None,
) -> list[Message]:
    """Normalize media parts by resolving local paths and validating references.

    This pass keeps remote URLs unchanged, converts local paths to absolute paths,
    and rejects unsupported ``data:`` URLs.
    """
    normalized_messages: list[Message] = []
    for msg in messages:
        parts: list[Part] = []
        for part in msg.content:
            if not isinstance(part, MediaPart):
                parts.append(part)
                continue

            resolved = resolve_image_ref(
                part.media.url,
                content_type=part.media.content_type,
                prompt_file_path=prompt_file_path,
                read_bytes=False,
            )
            if resolved.kind == 'remote_url':
                url = resolved.url or part.media.url
            else:
                url = str(resolved.file_path)
            media = MediaContent(url=url, content_type=resolved.mime_type)
            parts.append(part.model_copy(update={'media': media}))
        normalized_messages.append(msg.model_copy(update={'content': parts}))
    return normalized_messages
