from __future__ import annotations

import asyncio
import logging
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace

import pytest
from mcp import Client, ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from PIL import Image

from mcp_sd15.errors import ConfigError
from mcp_sd15.model_manager import ModelManager
from mcp_sd15.pipeline_factory import PipelineFactory
from mcp_sd15.schema import TextToImageRequest
from mcp_sd15.server import create_server
from mcp_sd15.service import ImageGenerationService


@pytest.fixture
def lazy_backend(app_config, monkeypatch):
    loads: list[int] = []
    calls: list[dict[str, object]] = []

    def pipeline(**kwargs):
        calls.append(kwargs)
        return SimpleNamespace(images=[Image.new("RGB", (4, 4), "red")])

    def create_base(self):
        loads.append(threading.get_ident())
        return pipeline

    monkeypatch.setattr(PipelineFactory, "create_base", create_base)
    backend = ModelManager(app_config)
    monkeypatch.setattr(backend, "_embed_prompts", lambda *args: ("positive", "negative"))
    monkeypatch.setattr(backend, "_generator", lambda seed: seed)

    def preprocess(image, **kwargs):
        assert backend.base is pipeline  # Base must load before any GPU/preprocessor work.
        return image

    monkeypatch.setattr(backend, "_depth", preprocess)
    monkeypatch.setattr(backend, "_openpose", preprocess)
    monkeypatch.setattr(backend, "_inpaint_condition", lambda image, mask: preprocess(image))
    monkeypatch.setattr(backend, "_pipeline_for", lambda *args, **kwargs: pipeline)
    return backend, loads, calls


async def test_discovery_does_not_load_model(app_config, lazy_backend) -> None:
    backend, loads, _ = lazy_backend
    assert backend.base is None
    server = create_server(ImageGenerationService(app_config, backend))
    assert loads == []
    async with Client(server, mode="legacy") as client:
        tools = (await client.list_tools()).tools
        assert len(tools) == 3
        assert loads == []
    assert backend.base is None


@pytest.mark.parametrize(
    ("tool", "arguments"),
    [
        ("generate_image", {}),
        ("generate_guided_image", {"control_mode": "depth"}),
        ("generate_guided_image", {"control_mode": "openpose"}),
        ("generate_guided_image", {"control_mode": "openpose_full"}),
        ("inpaint_image", {}),
    ],
)
async def test_first_generation_loads_once_in_worker(
    app_config, lazy_backend, image_base64, tool, arguments, caplog
) -> None:
    caplog.set_level(logging.INFO, logger="mcp_sd15.model_manager")
    backend, loads, calls = lazy_backend
    arguments = {"positive_prompt": "cat", **arguments}
    if tool != "generate_image":
        arguments["input_image"] = image_base64
    if tool == "inpaint_image":
        arguments["mask_image"] = image_base64
    server = create_server(ImageGenerationService(app_config, backend))
    async with Client(server) as client:
        for _ in range(2):
            result = await client.call_tool(tool, arguments)
            assert result.is_error is False
            assert result.content[0].type == "image"
    assert len(loads) == 1
    assert loads[0] != threading.get_ident()
    assert len(calls) == 2
    assert calls[0]["num_inference_steps"] == app_config.generation.steps
    assert calls[0]["guidance_scale"] == app_config.generation.cfg_scale
    assert caplog.text.count("model load started") == 1
    assert caplog.text.count("model load completed") == 1


async def test_simultaneous_requests_serialize_load_and_generation(
    app_config, lazy_backend
) -> None:
    backend, loads, calls = lazy_backend
    service = ImageGenerationService(app_config, backend)
    await asyncio.gather(
        *(service.generate_image(TextToImageRequest.create("cat")) for _ in range(4))
    )
    assert len(loads) == 1
    assert len(calls) == 4


def test_load_is_thread_safe_and_idempotent(app_config, monkeypatch) -> None:
    loads: list[int] = []
    pipeline = object()

    def create_base(self):
        loads.append(threading.get_ident())
        time.sleep(0.03)
        return pipeline

    monkeypatch.setattr(PipelineFactory, "create_base", create_base)
    backend = ModelManager(app_config)
    with ThreadPoolExecutor(max_workers=8) as executor:
        list(executor.map(lambda _: backend.load(), range(8)))
    backend.load()
    assert len(loads) == 1
    assert backend.base is pipeline


@pytest.mark.parametrize(
    ("tool", "arguments"),
    [
        ("generate_image", {"positive_prompt": ""}),
        ("generate_image", {"size": "invalid"}),
        ("generate_guided_image", {"control_mode": "invalid", "input_image": "invalid"}),
        ("generate_guided_image", {"control_mode": "depth", "input_image": "invalid"}),
        ("inpaint_image", {"input_image": "invalid", "mask_image": "invalid"}),
    ],
)
async def test_invalid_requests_do_not_load(app_config, lazy_backend, tool, arguments) -> None:
    backend, loads, _ = lazy_backend
    server = create_server(ImageGenerationService(app_config, backend))
    async with Client(server) as client:
        result = await client.call_tool(tool, {"positive_prompt": "cat", **arguments})
    assert result.is_error is True
    assert loads == []


async def test_invalid_mask_does_not_load(app_config, lazy_backend, image_base64) -> None:
    backend, loads, _ = lazy_backend
    async with Client(create_server(ImageGenerationService(app_config, backend))) as client:
        result = await client.call_tool(
            "inpaint_image",
            {"positive_prompt": "cat", "input_image": image_base64, "mask_image": "invalid"},
        )
    assert result.is_error is True
    assert loads == []


@pytest.mark.parametrize("tool", ["generate_guided_image", "inpaint_image"])
async def test_disabled_controlnet_does_not_load(app_config, lazy_backend, image_base64, tool):
    backend, loads, _ = lazy_backend
    config = replace(
        app_config,
        controlnet=replace(
            app_config.controlnet,
            depth=replace(app_config.controlnet.depth, enabled=False),
            inpaint=replace(app_config.controlnet.inpaint, enabled=False),
        ),
    )
    arguments = {"positive_prompt": "cat", "input_image": image_base64}
    if tool == "generate_guided_image":
        arguments["control_mode"] = "depth"
    else:
        arguments["mask_image"] = image_base64
    async with Client(create_server(ImageGenerationService(config, backend))) as client:
        result = await client.call_tool(tool, arguments)
    assert result.is_error is True
    assert loads == []


async def test_load_failure_is_tool_error_and_next_call_retries(
    app_config, lazy_backend, monkeypatch, caplog
):
    backend, loads, _ = lazy_backend
    create_base = backend.factory.create_base
    attempts = 0

    def fail_once():
        nonlocal attempts
        attempts += 1
        if attempts == 1:
            raise ConfigError("test model could not load")
        return create_base()

    monkeypatch.setattr(backend.factory, "create_base", fail_once)
    async with Client(create_server(ImageGenerationService(app_config, backend))) as client:
        result = await client.call_tool("generate_image", {"positive_prompt": "cat"})
        assert result.is_error is True
        assert "test model could not load" in caplog.text
        assert backend.base is None
        assert len((await client.list_tools()).tools) == 3
        result = await client.call_tool("generate_image", {"positive_prompt": "cat"})
        assert result.is_error is False
    assert attempts == 2
    assert len(loads) == 1


async def test_cli_stdio_discovery_forbids_inference_runtime(tmp_path: Path) -> None:
    checkpoint = tmp_path / "checkpoints" / "placeholder.safetensors"
    checkpoint.parent.mkdir()
    checkpoint.touch()
    config = tmp_path / "config.toml"
    config.write_text('[model]\nroot = "."\ncheckpoint = "placeholder.safetensors"\n')
    # Exercise the real CLI while making any inference import or model load fatal.
    bootstrap = """
import importlib.abc
import sys
class NoInference(importlib.abc.MetaPathFinder):
    def find_spec(self, fullname, path=None, target=None):
        if fullname.split('.')[0] in {'torch', 'diffusers', 'controlnet_aux', 'sd_embed'}:
            raise AssertionError('inference imported during discovery: ' + fullname)
sys.meta_path.insert(0, NoInference())
from mcp_sd15.pipeline_factory import PipelineFactory
def forbidden(self):
    raise AssertionError('model loaded during discovery')
PipelineFactory.create_base = forbidden
from mcp_sd15.cli import main
raise SystemExit(main())
"""
    params = StdioServerParameters(
        command=sys.executable,
        args=["-c", bootstrap, "--config", str(config)],
    )
    async with (
        stdio_client(params) as (read, write),
        ClientSession(read, write, read_timeout_seconds=15) as client,
    ):
        await client.initialize()
        tools = (await client.list_tools()).tools
        assert {tool.name for tool in tools} == {
            "generate_image", "generate_guided_image", "inpaint_image"
        }