from __future__ import annotations

import asyncio
import socket
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import AsyncMock
from urllib.error import URLError

import pytest
from pydantic_ai.messages import ImageUrl

from keble_helpers.ai.image_runtime import (
    ImagePreflightDecision,
    ImagePreflightReason,
    ImagePromptChecker,
    arun_with_image_fallback,
    extract_image_urls,
    is_image_url_model_404_error,
)


def test_extract_image_urls_reads_only_typed_image_parts() -> None:
    prompt = [
        "text-1",
        ImageUrl(url="https://img.example.com/1.jpg"),
        "text-2",
        ImageUrl(url="https://img.example.com/2.jpg"),
    ]
    assert extract_image_urls(prompt) == [
        "https://img.example.com/1.jpg",
        "https://img.example.com/2.jpg",
    ]


def test_preflight_status_policy_defaults_to_reject_404_only() -> None:
    runtime = ImagePromptChecker()
    rejected = runtime._decision_from_status_code(
        url="https://img.example.com/dead.jpg",
        status_code=404,
    )
    allowed = runtime._decision_from_status_code(
        url="https://img.example.com/redirect.jpg",
        status_code=301,
    )
    assert rejected.allowed is False
    assert rejected.reason == ImagePreflightReason.HTTP_404
    assert allowed.allowed is True
    assert allowed.reason == ImagePreflightReason.OK


def test_preflight_status_policy_only_accepts() -> None:
    runtime = ImagePromptChecker(only_accepts=[200, 206, 200])
    accepted = runtime._decision_from_status_code(
        url="https://img.example.com/ok.jpg",
        status_code=206,
    )
    rejected = runtime._decision_from_status_code(
        url="https://img.example.com/redirect.jpg",
        status_code=301,
    )
    assert accepted.allowed is True
    assert accepted.reason == ImagePreflightReason.OK
    assert rejected.allowed is False
    assert rejected.reason == ImagePreflightReason.HTTP_OTHER


def test_preflight_status_policy_rejects() -> None:
    runtime = ImagePromptChecker(rejects=[301, 404, 301])
    accepted = runtime._decision_from_status_code(
        url="https://img.example.com/ok.jpg",
        status_code=200,
    )
    rejected = runtime._decision_from_status_code(
        url="https://img.example.com/redirect.jpg",
        status_code=301,
    )
    assert accepted.allowed is True
    assert accepted.reason == ImagePreflightReason.OK
    assert rejected.allowed is False
    assert rejected.reason == ImagePreflightReason.HTTP_OTHER


def test_preflight_status_policy_rejects_mutual_configuration() -> None:
    with pytest.raises(
        ValueError,
        match="Provide only one status policy: `only_accepts` or `rejects`.",
    ):
        ImagePromptChecker(only_accepts=[200], rejects=[404])


def test_preflight_runtime_rejects_negative_image_404_retry_attempts() -> None:
    with pytest.raises(
        ValueError,
        match="image_model_404_retry_attempts must be >= 0",
    ):
        ImagePromptChecker(image_model_404_retry_attempts=-1)


def test_preflight_runtime_surfaces_redirect_status_without_following() -> None:
    class _RedirectHandler(BaseHTTPRequestHandler):
        def do_GET(self) -> None:  # noqa: N802
            if self.path == "/redirect":
                self.send_response(301)
                self.send_header("Location", "/ok")
                self.end_headers()
                return
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"ok")

        def log_message(self, format: str, *args) -> None:  # noqa: A003
            _ = (format, args)
            return

    with ThreadingHTTPServer(("127.0.0.1", 0), _RedirectHandler) as server:
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        port = int(server.server_address[1])
        redirect_url = f"http://127.0.0.1:{port}/redirect"

        permissive_runtime = ImagePromptChecker()
        permissive_decision = asyncio.run(
            permissive_runtime._aprobe_url(url=redirect_url)
        )
        strict_runtime = ImagePromptChecker(only_accepts=[200, 206])
        strict_decision = asyncio.run(strict_runtime._aprobe_url(url=redirect_url))

        server.shutdown()
        thread.join()

    assert permissive_decision.status_code == 301
    assert permissive_decision.allowed is True
    assert strict_decision.status_code == 301
    assert strict_decision.allowed is False


@pytest.mark.asyncio
async def test_preflight_keeps_only_ok_urls_and_drops_failures() -> None:
    runtime = ImagePromptChecker(
        enabled=True,
        timeout_secs=2.0,
        max_concurrency=4,
        cache_ttl_secs=600,
    )
    ok_url = "https://img.example.com/ok.jpg"
    dead_url = "https://img.example.com/dead.jpg"

    async def _probe(url: str) -> ImagePreflightDecision:
        if url == ok_url:
            return ImagePreflightDecision(
                url=url,
                allowed=True,
                reason=ImagePreflightReason.OK,
                status_code=200,
            )
        return ImagePreflightDecision(
            url=url,
            allowed=False,
            reason=ImagePreflightReason.HTTP_404,
            status_code=404,
        )

    probe_mock = AsyncMock(side_effect=_probe)
    runtime._aprobe_url = probe_mock  # type: ignore[method-assign]

    filtered_prompt, result = await runtime.afilter_prompt_images(
        prompt=[
            "input text",
            ImageUrl(url=dead_url),
            ImageUrl(url=ok_url),
            "tail text",
        ]
    )

    assert extract_image_urls(filtered_prompt) == [ok_url]
    assert result.kept_urls == [ok_url]
    assert [(item.url, item.reason) for item in result.dropped] == [
        (dead_url, ImagePreflightReason.HTTP_404)
    ]
    assert probe_mock.await_count == 2


@pytest.mark.asyncio
async def test_preflight_cache_hit_avoids_duplicate_probe() -> None:
    runtime = ImagePromptChecker(
        enabled=True,
        timeout_secs=2.0,
        max_concurrency=4,
        cache_ttl_secs=600,
    )
    url = "https://img.example.com/cache.jpg"
    probe_mock = AsyncMock(
        return_value=ImagePreflightDecision(
            url=url,
            allowed=True,
            reason=ImagePreflightReason.OK,
            status_code=200,
        )
    )
    runtime._aprobe_url = probe_mock  # type: ignore[method-assign]

    prompt = ["text", ImageUrl(url=url), ImageUrl(url=url)]
    filtered_prompt_1, _ = await runtime.afilter_prompt_images(prompt=prompt)
    filtered_prompt_2, _ = await runtime.afilter_prompt_images(prompt=prompt)

    assert extract_image_urls(filtered_prompt_1) == [url, url]
    assert extract_image_urls(filtered_prompt_2) == [url, url]
    assert probe_mock.await_count == 1


def test_is_image_url_model_404_error_matches_provider_message() -> None:
    assert is_image_url_model_404_error(
        RuntimeError("400 BadRequest: provided image url https://dead/a.png is 404")
    )
    assert not is_image_url_model_404_error(
        RuntimeError("400 BadRequest: invalid JSON payload")
    )


@pytest.mark.asyncio
async def test_arun_with_image_fallback_retries_once_text_only() -> None:
    dead_url = "https://img.example.com/dead.jpg"
    alive_url = "https://img.example.com/alive.jpg"

    class _StubAgent:
        def __init__(self) -> None:
            self.prompts: list[list[object]] = []

        async def run(self, prompt, **kwargs):  # type: ignore[no-untyped-def]
            _ = kwargs
            prompt_parts = list(prompt)
            self.prompts.append(prompt_parts)
            has_dead_url = any(
                isinstance(part, ImageUrl) and part.url == dead_url for part in prompt_parts
            )
            if has_dead_url:
                raise RuntimeError(
                    f"400 BadRequest: provided image url {dead_url} is 404"
                )
            return type("Result", (), {"output": "ok"})()

    agent = _StubAgent()
    runtime = ImagePromptChecker(enabled=False)
    result = await arun_with_image_fallback(
        agent=agent,  # type: ignore[arg-type]
        prompt=["input", ImageUrl(url=dead_url), ImageUrl(url=alive_url)],
        image_prompt_checker=runtime,
    )

    assert result.output == "ok"
    assert len(agent.prompts) == 2
    assert extract_image_urls(agent.prompts[0]) == [dead_url, alive_url]
    assert extract_image_urls(agent.prompts[1]) == []


@pytest.mark.asyncio
async def test_arun_with_image_fallback_reraise_preserves_exception() -> None:
    dead_url = "https://img.example.com/dead.jpg"

    class _AlwaysFailAgent:
        async def run(self, prompt, **kwargs):  # type: ignore[no-untyped-def]
            _ = (prompt, kwargs)
            raise RuntimeError(
                f"400 BadRequest: provided image url {dead_url} is 404"
            )

    with pytest.raises(RuntimeError, match="provided image url"):
        await arun_with_image_fallback(
            agent=_AlwaysFailAgent(),  # type: ignore[arg-type]
            prompt=["input", ImageUrl(url=dead_url)],
            image_prompt_checker=ImagePromptChecker(enabled=False),
        )


@pytest.mark.asyncio
async def test_runtime_method_arun_with_image_fallback_uses_instance_retry_policy() -> None:
    dead_url = "https://img.example.com/dead.jpg"

    class _StubAgent:
        def __init__(self) -> None:
            self.prompts: list[list[object]] = []

        async def run(self, prompt, **kwargs):  # type: ignore[no-untyped-def]
            _ = kwargs
            prompt_parts = list(prompt)
            self.prompts.append(prompt_parts)
            if len(self.prompts) == 1:
                raise RuntimeError(
                    f"400 BadRequest: provided image url {dead_url} is 404"
                )
            return type("Result", (), {"output": "ok"})()

    agent = _StubAgent()
    runtime = ImagePromptChecker(
        enabled=False,
        image_model_404_retry_attempts=1,
    )

    result = await runtime.arun_with_image_fallback(
        agent=agent,  # type: ignore[arg-type]
        prompt=["input", ImageUrl(url=dead_url)],
    )

    assert result.output == "ok"
    assert len(agent.prompts) == 2
    assert extract_image_urls(agent.prompts[0]) == [dead_url]
    assert extract_image_urls(agent.prompts[1]) == []


@pytest.mark.asyncio
async def test_arun_with_image_fallback_reuses_checker_retry_policy_without_prefilter() -> None:
    dead_url = "https://img.example.com/dead.jpg"

    class _StubAgent:
        def __init__(self) -> None:
            self.prompts: list[list[object]] = []

        async def run(self, prompt, **kwargs):  # type: ignore[no-untyped-def]
            _ = kwargs
            prompt_parts = list(prompt)
            self.prompts.append(prompt_parts)
            if len(self.prompts) == 1:
                raise RuntimeError(
                    f"400 BadRequest: provided image url {dead_url} is 404"
                )
            return type("Result", (), {"output": "ok"})()

    agent = _StubAgent()
    checker = ImagePromptChecker(enabled=False, image_model_404_retry_attempts=1)

    result = await arun_with_image_fallback(
        agent=agent,  # type: ignore[arg-type]
        prompt=["input", ImageUrl(url=dead_url)],
        image_prompt_checker=checker,
        preflight_prompt=False,
    )

    assert result.output == "ok"
    assert len(agent.prompts) == 2
    assert extract_image_urls(agent.prompts[0]) == [dead_url]
    assert extract_image_urls(agent.prompts[1]) == []


def test_preflight_decision_maps_timeout_network() -> None:
    runtime = ImagePromptChecker()
    timeout_decision = runtime._decision_from_exception(
        url="https://img.example.com/a.jpg",
        exception=socket.timeout("timeout"),
    )
    network_decision = runtime._decision_from_exception(
        url="https://img.example.com/b.jpg",
        exception=URLError("network"),
    )

    assert timeout_decision.reason == ImagePreflightReason.TIMEOUT
    assert network_decision.reason == ImagePreflightReason.NETWORK


def test_preflight_runtime_is_loop_safe_for_reuse() -> None:
    runtime = ImagePromptChecker(
        enabled=True,
        timeout_secs=2.0,
        max_concurrency=4,
        cache_ttl_secs=600,
    )
    runtime._blocking_fetch_status_code = lambda *_: 200  # type: ignore[method-assign]
    prompt = ["input text", ImageUrl(url="https://img.example.com/loop-safe.jpg")]

    async def _run_once() -> None:
        filtered_prompt, result = await runtime.afilter_prompt_images(prompt=prompt)
        assert extract_image_urls(filtered_prompt) == [
            "https://img.example.com/loop-safe.jpg"
        ]
        assert result.kept_urls == ["https://img.example.com/loop-safe.jpg"]

    asyncio.run(_run_once())
    asyncio.run(_run_once())


@pytest.mark.asyncio
async def test_preflight_runtime_evicts_expired_cache_entries() -> None:
    runtime = ImagePromptChecker(
        enabled=True,
        timeout_secs=2.0,
        max_concurrency=1,
        cache_ttl_secs=1,
    )
    runtime._cache["https://img.example.com/expired.jpg"] = (
        0.0,
        ImagePreflightDecision(
            url="https://img.example.com/expired.jpg",
            allowed=False,
            reason=ImagePreflightReason.HTTP_404,
            status_code=404,
        ),
    )
    runtime._aprobe_url = AsyncMock(
        return_value=ImagePreflightDecision(
            url="https://img.example.com/alive.jpg",
            allowed=True,
            reason=ImagePreflightReason.OK,
            status_code=200,
        )
    )  # type: ignore[method-assign]

    await runtime._aget_or_probe(url="https://img.example.com/alive.jpg")

    assert "https://img.example.com/expired.jpg" not in runtime._cache
