"""Validation utilities for image sources (URLs and local file paths).

Prevents SSRF (Server-Side Request Forgery) by:
- Restricting URL schemes to http/https
- Resolving the hostname and rejecting private/loopback/link-local/reserved/
  multicast IPs
- Restricting file path access to known image extensions

Also exposes a redirect-aware fetcher (``safe_fetch``) that re-validates each
intermediate URL hop, so an attacker cannot use an open redirect to reach
internal addresses.
"""

import ipaddress
import socket
from pathlib import Path
from typing import Iterable
from urllib.parse import urlparse

import httpx
from fastapi import HTTPException


# Image file extensions considered safe to serve
_IMAGE_EXTENSIONS = frozenset(
    {
        ".jpg",
        ".jpeg",
        ".png",
        ".bmp",
        ".gif",
        ".webp",
        ".tiff",
        ".tif",
        ".ico",
    }
)


def _is_blocked_ip(ip: str) -> bool:
    """Return True when *ip* belongs to a category we refuse to fetch from."""
    try:
        addr = ipaddress.ip_address(ip)
    except ValueError:
        return True  # unparseable → block
    return (
        addr.is_private
        or addr.is_loopback
        or addr.is_link_local
        or addr.is_reserved
        or addr.is_multicast
        or addr.is_unspecified
    )


def _resolve_hostname(hostname: str) -> Iterable[str]:
    """Resolve *hostname* to all IPv4/IPv6 addresses."""
    try:
        infos = socket.getaddrinfo(hostname, None)
    except socket.gaierror as exc:
        raise HTTPException(
            status_code=400, detail=f"Cannot resolve hostname: {hostname} ({exc})"
        ) from exc
    return {info[4][0] for info in infos}


def validate_image_url(url: str) -> None:
    """Validate that *url* is safe to fetch.

    Checks:
    - Scheme is http/https
    - Hostname is present
    - Hostname is not a literal blocked IP
    - DNS resolves to non-private, non-loopback, non-link-local, non-reserved,
      non-multicast addresses (SSRF protection)
    """
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported URL scheme: {parsed.scheme!r}. Only http and https are allowed.",
        )

    hostname = parsed.hostname
    if not hostname:
        raise HTTPException(status_code=400, detail="URL has no hostname")

    # If hostname IS an IP literal, skip DNS but still check the category.
    try:
        ipaddress.ip_address(hostname)
        ips: Iterable[str] = [hostname]
    except ValueError:
        ips = _resolve_hostname(hostname)

    for ip in ips:
        if _is_blocked_ip(ip):
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Refusing to fetch URL: hostname {hostname!r} resolves to "
                    f"blocked address {ip} (private / loopback / link-local / "
                    "reserved / multicast)"
                ),
            )


def validate_image_path(file_path: str | Path) -> Path:
    """Validate a local file path points to a real image file.

    Checks:
    - The extension is a known image format
    - The resolved path does not escape via symlinks to unexpected locations

    Returns the resolved Path on success, raises HTTPException on violation.
    """
    resolved = Path(file_path).resolve()

    suffix = resolved.suffix.lower()
    if suffix not in _IMAGE_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type: {suffix!r}. Only image files are allowed.",
        )

    return resolved


# ---------------------------------------------------------------------------
# Safe HTTP fetch (manual redirect handling with per-hop validation)
# ---------------------------------------------------------------------------


_DEFAULT_MAX_HOPS = 3


async def safe_fetch(
    url: str,
    *,
    timeout: float = 15.0,
    max_hops: int = _DEFAULT_MAX_HOPS,
) -> httpx.Response:
    """Fetch *url* with SSRF protection and manual redirect handling.

    Each redirect target is re-validated via :func:`validate_image_url` so
    an attacker cannot use an open redirect at a public host to pivot to
    internal addresses.

    Raises:
        HTTPException: on validation failure or too many redirects.
        httpx.HTTPStatusError / httpx.RequestError: for transport errors.
    """
    current = url
    async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client:
        for _ in range(max_hops + 1):
            validate_image_url(current)
            response = await client.get(current)
            if response.is_redirect:
                location = response.headers.get("location")
                if not location:
                    raise HTTPException(status_code=400, detail="Redirect without Location header")
                # Resolve relative redirects against the current URL
                next_url = str(httpx.URL(current).join(location))
                current = next_url
                continue
            response.raise_for_status()
            return response
    raise HTTPException(status_code=400, detail=f"Too many redirects (max {max_hops})")
