from __future__ import annotations

import asyncio
import socket
import time
from collections.abc import Sequence
from enum import Enum
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import HTTPRedirectHandler, Request, build_opener

from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.messages import ImageUrl, UserContent
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential

from ..typings import PydanticModelConfig


class ImagePreflightReason(str, Enum):
    OK = "ok"
    HTTP_404 = "http_404"
    HTTP_403 = "http_403"
    HTTP_OTHER = "http_other"
    TIMEOUT = "timeout"
    NETWORK = "network"


class ImagePreflightDecision(BaseModel):
    """Preflight decision for one image URL used in prompt parts."""

    model_config = PydanticModelConfig.default()

    url: str = Field(min_length=1)
    allowed: bool
    reason: ImagePreflightReason
    status_code: int | None = Field(default=None)


class ImagePreflightBatchResult(BaseModel):
    """Batch preflight result for one prompt assembly."""

    model_config = PydanticModelConfig.default()

    kept_urls: list[str] = Field(default_factory=list)
    dropped: list[ImagePreflightDecision] = Field(default_factory=list)


def extract_image_urls(prompt: Sequence[UserContent]) -> list[str]:
    return [part.url for part in prompt if isinstance(part, ImageUrl)]


def strip_image_url_parts(prompt: Sequence[UserContent]) -> list[UserContent]:
    return [part for part in prompt if not isinstance(part, ImageUrl)]


def is_image_url_model_404_error(exception: BaseException) -> bool:
    normalized = str(exception).casefold()
    if "404" not in normalized:
        return False
    if "image url" not in normalized and "image_url" not in normalized:
        return False
    return any(
        token in normalized
        for token in (
            "provided image",
            "inaccessible image",
            "url is not accessible",
            "not reachable",
            "not found",
        )
    )


class ImagePromptPreflightRuntime:
    """Runtime image URL preflight with bounded concurrency and short-lived cache."""

    def __init__(
        self,
        *,
        enabled: bool = True,
        timeout_secs: float = 2.0,
        max_concurrency: int = 4,
        cache_ttl_secs: int = 600,
        only_accepts: list[int] | None = None,
        rejects: list[int] | None = None,
    ) -> None:
        if timeout_secs <= 0:
            raise ValueError("timeout_secs must be > 0")
        if max_concurrency < 1:
            raise ValueError("max_concurrency must be >= 1")
        if cache_ttl_secs < 1:
            raise ValueError("cache_ttl_secs must be >= 1")
        if only_accepts is not None and rejects is not None:
            raise ValueError(
                "Provide only one status policy: `only_accepts` or `rejects`."
            )

        self.enabled = enabled
        self.timeout_secs = timeout_secs
        self.max_concurrency = max_concurrency
        self.cache_ttl_secs = cache_ttl_secs
        self._only_accepts = self._normalize_status_codes(
            statuses=only_accepts,
            field_name="only_accepts",
        )
        self._rejects = self._normalize_status_codes(
            statuses=rejects,
            field_name="rejects",
        )
        if self._only_accepts is None and self._rejects is None:
            self._rejects = (404,)
        self._semaphore: asyncio.Semaphore | None = None
        self._semaphore_loop_key: int | None = None
        self._cache_max_entries: int = max(1024, max_concurrency * 1024)
        self._cache: dict[str, tuple[float, ImagePreflightDecision]] = {}

    async def afilter_prompt_images(
        self,
        *,
        prompt: Sequence[UserContent],
    ) -> tuple[list[UserContent], ImagePreflightBatchResult]:
        prompt_parts = list(prompt)
        if not prompt_parts:
            return [], ImagePreflightBatchResult()
        if not self.enabled:
            return (
                prompt_parts,
                ImagePreflightBatchResult(kept_urls=extract_image_urls(prompt_parts)),
            )

        decisions_by_url = await self._apreflight_urls(urls=extract_image_urls(prompt_parts))
        filtered_prompt: list[UserContent] = []
        kept_urls: list[str] = []
        dropped: list[ImagePreflightDecision] = []
        for part in prompt_parts:
            if not isinstance(part, ImageUrl):
                filtered_prompt.append(part)
                continue
            decision = decisions_by_url.get(part.url)
            if decision is None or decision.allowed:
                filtered_prompt.append(part)
                kept_urls.append(part.url)
                continue
            dropped.append(decision)

        return filtered_prompt, ImagePreflightBatchResult(
            kept_urls=kept_urls,
            dropped=dropped,
        )

    async def _apreflight_urls(
        self,
        *,
        urls: Sequence[str],
    ) -> dict[str, ImagePreflightDecision]:
        deduped_urls: list[str] = []
        seen: set[str] = set()
        for url in urls:
            normalized = url.strip()
            if not normalized or normalized in seen:
                continue
            seen.add(normalized)
            deduped_urls.append(normalized)
        if not deduped_urls:
            return {}

        decisions = await asyncio.gather(*[self._aget_or_probe(url=url) for url in deduped_urls])
        return {decision.url: decision for decision in decisions}

    async def _aget_or_probe(self, *, url: str) -> ImagePreflightDecision:
        now = time.monotonic()
        self._evict_expired_cache(now=now)
        cached = self._cache.get(url)
        if cached is not None and cached[0] >= now:
            return cached[1]
        decision = await self._aprobe_url(url=url)
        if len(self._cache) >= self._cache_max_entries:
            self._evict_oldest_cache_entry()
        self._cache[url] = (now + float(self.cache_ttl_secs), decision)
        return decision

    async def _aprobe_url(self, *, url: str) -> ImagePreflightDecision:
        if not url.lower().startswith(("http://", "https://")):
            return ImagePreflightDecision(
                url=url,
                allowed=False,
                reason=ImagePreflightReason.NETWORK,
                status_code=None,
            )

        async with self._get_semaphore():
            probe_result = (
                await asyncio.gather(
                    asyncio.to_thread(
                        self._blocking_fetch_status_code,
                        url,
                        self.timeout_secs,
                    ),
                    return_exceptions=True,
                )
            )[0]
        if isinstance(probe_result, BaseException) and not isinstance(probe_result, Exception):
            raise probe_result
        if isinstance(probe_result, Exception):
            return self._decision_from_exception(url=url, exception=probe_result)

        return self._decision_from_status_code(url=url, status_code=int(probe_result))

    def _get_semaphore(self) -> asyncio.Semaphore:
        # Runtime instances can be reused by different event loops (for example server reload/test loops).
        # Keep a loop-local semaphore to avoid "Future attached to a different loop" failures.
        loop_key = id(asyncio.get_running_loop())
        if self._semaphore is not None and self._semaphore_loop_key == loop_key:
            return self._semaphore
        self._semaphore = asyncio.Semaphore(self.max_concurrency)
        self._semaphore_loop_key = loop_key
        return self._semaphore

    def _evict_expired_cache(self, *, now: float) -> None:
        expired_urls = [
            cached_url
            for cached_url, (expires_at, _) in self._cache.items()
            if expires_at < now
        ]
        for cached_url in expired_urls:
            self._cache.pop(cached_url, None)

    def _evict_oldest_cache_entry(self) -> None:
        if not self._cache:
            return
        oldest_url = min(
            self._cache.items(),
            key=lambda item: item[1][0],
        )[0]
        self._cache.pop(oldest_url, None)

    @staticmethod
    def _normalize_status_codes(
        *,
        statuses: list[int] | None,
        field_name: str,
    ) -> tuple[int, ...] | None:
        if statuses is None:
            return None
        if len(statuses) == 0:
            raise ValueError(f"{field_name} must not be empty.")
        normalized: list[int] = []
        seen: set[int] = set()
        for status_code in statuses:
            if isinstance(status_code, bool) or not isinstance(status_code, int):
                raise ValueError(
                    f"{field_name} must contain integer HTTP status codes."
                )
            if status_code < 100 or status_code > 599:
                raise ValueError(
                    f"{field_name} values must be valid HTTP status codes (100-599)."
                )
            if status_code in seen:
                continue
            seen.add(status_code)
            normalized.append(status_code)
        return tuple(normalized)

    @staticmethod
    def _blocking_fetch_status_code(url: str, timeout_secs: float) -> int:
        class _NoRedirectHandler(HTTPRedirectHandler):
            def redirect_request(self, req, fp, code, msg, headers, newurl):  # type: ignore[override]
                _ = (req, fp, code, msg, headers, newurl)
                return None

        opener = build_opener(_NoRedirectHandler)
        request = Request(
            url=url,
            method="GET",
            headers={
                "Range": "bytes=0-0",
                "User-Agent": "keble-helpers-image-preflight/1.0",
            },
        )
        with opener.open(request, timeout=timeout_secs) as response:
            return int(response.status)

    def _decision_from_status_code(
        self,
        *,
        url: str,
        status_code: int,
    ) -> ImagePreflightDecision:
        if self._is_status_allowed(status_code=status_code):
            return ImagePreflightDecision(
                url=url,
                allowed=True,
                reason=ImagePreflightReason.OK,
                status_code=status_code,
            )
        if status_code == 404:
            reason = ImagePreflightReason.HTTP_404
        elif status_code == 403:
            reason = ImagePreflightReason.HTTP_403
        else:
            reason = ImagePreflightReason.HTTP_OTHER
        return ImagePreflightDecision(
            url=url,
            allowed=False,
            reason=reason,
            status_code=status_code,
        )

    def _decision_from_exception(
        self,
        *,
        url: str,
        exception: Exception,
    ) -> ImagePreflightDecision:
        if isinstance(exception, HTTPError):
            return self._decision_from_status_code(
                url=url,
                status_code=int(exception.code),
            )
        if isinstance(exception, (TimeoutError, socket.timeout)):
            return ImagePreflightDecision(
                url=url,
                allowed=False,
                reason=ImagePreflightReason.TIMEOUT,
                status_code=None,
            )
        if isinstance(exception, URLError):
            if isinstance(exception.reason, (TimeoutError, socket.timeout)):
                return ImagePreflightDecision(
                    url=url,
                    allowed=False,
                    reason=ImagePreflightReason.TIMEOUT,
                    status_code=None,
                )
            return ImagePreflightDecision(
                url=url,
                allowed=False,
                reason=ImagePreflightReason.NETWORK,
                status_code=None,
            )
        return ImagePreflightDecision(
            url=url,
            allowed=False,
            reason=ImagePreflightReason.NETWORK,
            status_code=None,
        )

    def _is_status_allowed(self, *, status_code: int) -> bool:
        if self._only_accepts is not None:
            return status_code in self._only_accepts
        if self._rejects is not None:
            return status_code not in self._rejects
        return True


async def arun_with_image_fallback(
    *,
    agent: Agent[Any, Any],
    prompt: str | Sequence[UserContent],
    preflight: ImagePromptPreflightRuntime | None = None,
    model_404_retry_attempts: int = 1,
    **agent_run_kwargs: Any,
) -> Any:
    if model_404_retry_attempts < 0:
        raise ValueError("model_404_retry_attempts must be >= 0")

    if isinstance(prompt, str):
        return await agent.run(prompt, **agent_run_kwargs)

    prompt_parts: list[UserContent] = list(prompt)
    if preflight is not None:
        prompt_parts, _ = await preflight.afilter_prompt_images(prompt=prompt_parts)

    prompt_holder: dict[str, list[UserContent]] = {"prompt": prompt_parts}
    max_attempts = model_404_retry_attempts + 1

    def _before_sleep(retry_state: Any) -> None:
        outcome = retry_state.outcome
        if outcome is None or not outcome.failed:
            return
        exception = outcome.exception()
        if exception is None:
            return
        if is_image_url_model_404_error(exception):
            prompt_holder["prompt"] = strip_image_url_parts(prompt_holder["prompt"])

    @retry(
        stop=stop_after_attempt(max_attempts),
        wait=wait_exponential(multiplier=0.2, min=0.2, max=1.0),
        retry=retry_if_exception(is_image_url_model_404_error),
        reraise=True,
        before_sleep=_before_sleep,
    )
    async def _run() -> Any:
        return await agent.run(prompt_holder["prompt"], **agent_run_kwargs)

    return await _run()
