"""End-to-end integration tests for Scheduler IR components."""

from __future__ import annotations

import json
import threading
from datetime import datetime, timedelta
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from types import SimpleNamespace

import pytest

from sagellm_control.engine_client import (
    ClientConfig,
    EngineClient,
    RequestContext,
    RetryPolicy,
)
from sagellm_control.ir.builder import IRBuilder
from sagellm_control.ir.executor import DefaultIRExecutor
from sagellm_control.ir.optimizer import IROptimizer
from sagellm_control.local_engine_client import LocalEngineClient
from sagellm_control.policies import FIFOPolicy, PriorityPolicy, SLOAwarePolicy
from sagellm_control.types import EngineInfo, EngineState, RequestMetadata, RequestPriority
from sagellm_protocol import Metrics, Request, Response


class _InMemoryLocalEngine:
    def __init__(self, engine_id: str) -> None:
        self.engine_id = engine_id

    async def execute(self, request: Request) -> Response:
        return Response(
            request_id=request.request_id,
            trace_id=request.trace_id,
            output_text=f"local-{request.request_id}",
            output_tokens=[1, 2],
            finish_reason="stop",
            metrics=Metrics(
                ttft_ms=5.0,
                tbt_ms=1.0,
                tpot_ms=1.0,
                throughput_tps=10.0,
                peak_mem_mb=64,
                error_rate=0.0,
            ),
        )

    async def stream(self, request: Request):
        yield SimpleNamespace(
            event="start", request_id=request.request_id, trace_id=request.trace_id
        )
        yield SimpleNamespace(
            event="end",
            request_id=request.request_id,
            trace_id=request.trace_id,
            content="done",
            output_tokens=[1],
            finish_reason="stop",
            metrics=Metrics(
                ttft_ms=1.0,
                tbt_ms=1.0,
                tpot_ms=1.0,
                throughput_tps=1.0,
                peak_mem_mb=1,
                error_rate=0.0,
            ),
        )

    async def health_check(self) -> bool:
        return True


class _LocalClientAdapter:
    def __init__(self, local_client: LocalEngineClient) -> None:
        self._local_client = local_client

    async def execute(self, request: Request) -> Response:
        return await self._local_client.execute_request(request)


class _HTTPClientAdapter:
    def __init__(self, client: EngineClient, engine: EngineInfo) -> None:
        self._client = client
        self._engine = engine

    async def execute(self, request: Request) -> SimpleNamespace:
        data = {
            "request_id": request.request_id,
            "trace_id": request.trace_id,
            "model": request.model,
            "prompt": request.prompt,
            "max_tokens": request.max_tokens,
            "stream": False,
        }
        result = await self._client.execute(
            self._engine,
            data,
            RequestContext(request_id=request.request_id, trace_id=request.trace_id),
        )
        return SimpleNamespace(
            request_id=result["request_id"],
            output_text=result.get("output_text", ""),
            finish_reason=result.get("finish_reason", "stop"),
        )


class _HTTPTestState:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.attempts = 0
        self.failures_before_success = 0


def _build_http_handler(state: _HTTPTestState):
    class _Handler(BaseHTTPRequestHandler):
        def _send_json(self, status_code: int, payload: dict) -> None:
            body = json.dumps(payload).encode("utf-8")
            self.send_response(status_code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

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

        def do_GET(self) -> None:  # noqa: N802
            if self.path == "/health":
                self._send_json(200, {"status": "ok"})
            else:
                self._send_json(404, {"error": "not found"})

        def do_POST(self) -> None:  # noqa: N802
            if self.path != "/v1/completions":
                self._send_json(404, {"error": "not found"})
                return

            content_length = int(self.headers.get("Content-Length", "0"))
            raw = self.rfile.read(content_length) if content_length > 0 else b"{}"
            request_data = json.loads(raw.decode("utf-8"))

            with state.lock:
                state.attempts += 1
                should_fail = state.attempts <= state.failures_before_success

            if should_fail:
                self._send_json(503, {"error": "temporary unavailable"})
                return

            request_id = request_data.get("request_id", "unknown")
            self._send_json(
                200,
                {
                    "request_id": request_id,
                    "output_text": f"http-{request_id}",
                    "finish_reason": "stop",
                    "metrics": {
                        "ttft_ms": 5.0,
                        "tbt_ms": 1.0,
                        "tpot_ms": 1.0,
                        "throughput_tps": 10.0,
                        "peak_mem_mb": 64,
                        "error_rate": 0.0,
                    },
                },
            )

    return _Handler


@pytest.fixture
def http_test_server() -> tuple[str, int, _HTTPTestState]:
    state = _HTTPTestState()
    server = ThreadingHTTPServer(("127.0.0.1", 0), _build_http_handler(state))
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    host, port = server.server_address
    try:
        yield host, port, state
    finally:
        server.shutdown()
        server.server_close()
        thread.join(timeout=2.0)


def _sample_request_metadata(request_id: str, priority: RequestPriority) -> RequestMetadata:
    return RequestMetadata(
        request_id=request_id,
        trace_id=f"trace-{request_id}",
        model_id="Qwen2-7B",
        prompt="hello integration",
        max_tokens=16,
        prompt_tokens=2,
        priority=priority,
    )


def test_ir_with_fifo_policy() -> None:
    """IR should reflect FIFO-selected request order."""
    policy = FIFOPolicy()
    req1 = _sample_request_metadata("req-fifo-1", RequestPriority.LOW)
    req2 = _sample_request_metadata("req-fifo-2", RequestPriority.CRITICAL)
    policy.add_request(req1)
    policy.add_request(req2)

    selected = policy.next_request()
    assert selected is not None
    builder = IRBuilder()
    ir = builder.build_from_request(selected)
    task = ir.get_node("task-req-fifo-1")
    assert task is not None
    assert task.request_id == "req-fifo-1"


def test_ir_with_priority_policy() -> None:
    """IR should reflect priority policy decision."""
    policy = PriorityPolicy()
    req_low = _sample_request_metadata("req-prio-low", RequestPriority.LOW)
    req_high = _sample_request_metadata("req-prio-critical", RequestPriority.CRITICAL)
    policy.add_request(req_low)
    policy.add_request(req_high)

    selected = policy.next_request()
    assert selected is not None
    builder = IRBuilder()
    ir = builder.build_from_request(selected)
    task = ir.get_node("task-req-prio-critical")
    assert task is not None
    assert task.priority == RequestPriority.CRITICAL.value
    assert ir.optimization_hints["priority"] == RequestPriority.CRITICAL.value


def test_ir_with_slo_aware_policy() -> None:
    """IR should reflect SLO-aware urgency decision."""
    policy = SLOAwarePolicy()
    now = datetime.now()
    req_loose = RequestMetadata(
        request_id="req-slo-loose",
        trace_id="trace-slo-loose",
        model_id="Qwen2-7B",
        prompt="loose",
        max_tokens=8,
        prompt_tokens=1,
        priority=RequestPriority.NORMAL,
        slo_deadline_ms=5000,
        arrival_time=now,
    )
    req_tight = RequestMetadata(
        request_id="req-slo-tight",
        trace_id="trace-slo-tight",
        model_id="Qwen2-7B",
        prompt="tight",
        max_tokens=8,
        prompt_tokens=1,
        priority=RequestPriority.NORMAL,
        slo_deadline_ms=100,
        arrival_time=now - timedelta(milliseconds=50),
    )
    policy.add_request(req_loose)
    policy.add_request(req_tight)

    selected = policy.next_request()
    assert selected is not None
    builder = IRBuilder()
    ir = builder.build_from_request(selected)
    task = ir.get_node("task-req-slo-tight")
    assert task is not None
    assert task.request_id == "req-slo-tight"


@pytest.mark.asyncio
async def test_ir_executor_via_local_engine_client() -> None:
    """IR executor should execute through LocalEngineClient adapter."""
    local_engine = _InMemoryLocalEngine(engine_id="local-1")
    local_client = LocalEngineClient(local_engine)
    adapter = _LocalClientAdapter(local_client)
    executor = DefaultIRExecutor(engine_client=adapter)

    builder = IRBuilder()
    optimizer = IROptimizer()
    metadata = _sample_request_metadata("req-local", RequestPriority.NORMAL)
    ir = optimizer.optimize(builder.build_from_request(metadata))

    results = await executor.execute(ir)

    assert "cmd-req-local" in results
    assert results["cmd-req-local"]["output_text"] == "local-req-local"


@pytest.mark.asyncio
async def test_ir_executor_via_http_engine_client(
    http_test_server: tuple[str, int, _HTTPTestState],
) -> None:
    """IR executor should execute through EngineClient HTTP adapter."""
    host, port, _ = http_test_server

    engine = EngineInfo(
        engine_id="http-1",
        model_id="Qwen2-7B",
        host=host,
        port=port,
        state=EngineState.READY,
    )
    http_client = EngineClient(ClientConfig(max_retries=0))
    adapter = _HTTPClientAdapter(http_client, engine)
    executor = DefaultIRExecutor(engine_client=adapter)

    builder = IRBuilder()
    ir = builder.build_from_request(_sample_request_metadata("req-http", RequestPriority.NORMAL))
    results = await executor.execute(ir)

    assert "cmd-req-http" in results
    assert results["cmd-req-http"]["output_text"] == "http-req-http"


@pytest.mark.asyncio
async def test_engine_client_retry_and_failure_paths(
    http_test_server: tuple[str, int, _HTTPTestState],
) -> None:
    """EngineClient should retry transient failures and surface final failure."""
    host, port, state = http_test_server
    state.failures_before_success = 2

    engine = EngineInfo(
        engine_id="retry-1",
        model_id="Qwen2-7B",
        host=host,
        port=port,
        state=EngineState.READY,
    )

    client = EngineClient(
        ClientConfig(
            max_retries=2,
            retry_policy=RetryPolicy.NONE,
        )
    )
    result = await client.execute(
        engine,
        {
            "request_id": "req-retry",
            "trace_id": "trace-retry",
            "model": "Qwen2-7B",
            "prompt": "retry",
            "max_tokens": 8,
            "stream": False,
        },
        RequestContext(request_id="req-retry", trace_id="trace-retry"),
    )

    assert state.attempts == 3
    assert result["output_text"] == "http-req-retry"

    state.failures_before_success = 100

    builder = IRBuilder()
    ir = builder.build_from_request(_sample_request_metadata("req-fail", RequestPriority.NORMAL))
    adapter = _HTTPClientAdapter(client, engine)
    executor = DefaultIRExecutor(engine_client=adapter)
    results = await executor.execute(ir)

    assert "cmd-req-fail" in results
    assert "error" in results["cmd-req-fail"]
    assert "temporary unavailable" in results["cmd-req-fail"]["error"]
