"""Integration tests for IR with ControlPlaneManager."""

from __future__ import annotations

import asyncio
import json
import threading
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import pytest

from sagellm_control import ControlPlaneManager, EngineState
from sagellm_control.ir.builder import IRBuilder
from sagellm_control.ir.executor import DefaultIRExecutor
from sagellm_control.ir.optimizer import IROptimizer
from sagellm_control.types import RequestMetadata, RequestPriority


def _register_ready_engine(manager: ControlPlaneManager, engine_id: str = "engine-1") -> None:
    manager.register_engine(
        engine_id=engine_id,
        model_id="Qwen2-7B",
        host="127.0.0.1",
        port=8000,
    )
    manager.update_engine_state(engine_id, EngineState.READY)


class _HTTPTestState:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.attempts = 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_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

            request_id = request_data.get("request_id", "unknown")
            self._send_json(
                200,
                {
                    "request_id": request_id,
                    "trace_id": request_data.get("trace_id", ""),
                    "output_text": f"cp-http-{request_id}",
                    "output_tokens": [1, 2],
                    "finish_reason": "stop",
                    "metrics": {
                        "ttft_ms": 12.0,
                        "tbt_ms": 1.0,
                        "tpot_ms": 1.0,
                        "throughput_tps": 40.0,
                        "peak_mem_mb": 128,
                        "error_rate": 0.0,
                    },
                },
            )

    return _Handler


@pytest.fixture
def http_engine_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)


@pytest.mark.asyncio
async def test_request_ir_optimize_execute_full_flow(
    http_engine_server: tuple[str, int, _HTTPTestState],
) -> None:
    """Should run request -> IR -> optimize -> execute using control plane backend."""
    host, port, state = http_engine_server
    manager = ControlPlaneManager()
    manager.register_engine(
        engine_id="engine-1",
        model_id="Qwen2-7B",
        host=host,
        port=port,
    )
    manager.update_engine_state("engine-1", EngineState.READY)

    metadata = RequestMetadata(
        request_id="req-flow-1",
        trace_id="trace-flow-1",
        model_id="Qwen2-7B",
        prompt="hello integration flow",
        max_tokens=16,
        prompt_tokens=3,
        priority=RequestPriority.NORMAL,
    )

    builder = IRBuilder(default_tp_degree=2)
    optimizer = IROptimizer()
    executor = DefaultIRExecutor(control_plane=manager)

    ir = builder.build_from_request(metadata)
    optimized_ir = optimizer.optimize(ir)
    results = await executor.execute(optimized_ir)

    assert "cmd-req-flow-1" in results
    result = results["cmd-req-flow-1"]
    assert result["request_id"] == "req-flow-1"
    assert result["output_text"] == "cp-http-req-flow-1"
    assert result["finish_reason"] == "stop"
    assert manager.get_engine("engine-1").active_requests == 0
    assert state.attempts >= 1


@pytest.mark.asyncio
async def test_multi_request_concurrent_scheduling() -> None:
    """Should schedule concurrent requests without failures."""
    manager = ControlPlaneManager()
    _register_ready_engine(manager, "engine-1")
    _register_ready_engine(manager, "engine-2")

    async def _schedule(idx: int):
        return await manager.schedule_request(
            request_id=f"req-{idx}",
            trace_id=f"trace-{idx}",
            model_id="Qwen2-7B",
            prompt="concurrent request",
            max_tokens=16,
        )

    decisions = await asyncio.gather(*[_schedule(i) for i in range(20)])

    assert all(decision.is_scheduled for decision in decisions)
    assert {decision.engine_id for decision in decisions} <= {"engine-1", "engine-2"}
    assert manager.get_engine("engine-1").active_requests >= 0
    assert manager.get_engine("engine-2").active_requests >= 0


@pytest.mark.asyncio
async def test_batch_ir_execution_with_mixed_priorities(
    http_engine_server: tuple[str, int, _HTTPTestState],
) -> None:
    """Should execute batch IR and preserve per-request outputs."""
    host, port, _ = http_engine_server
    manager = ControlPlaneManager()
    manager.register_engine(
        engine_id="engine-1",
        model_id="Qwen2-7B",
        host=host,
        port=port,
    )
    manager.update_engine_state("engine-1", EngineState.READY)

    now = datetime.now()
    requests = [
        RequestMetadata(
            request_id="req-batch-critical",
            trace_id="trace-batch-critical",
            model_id="Qwen2-7B",
            prompt="critical",
            max_tokens=8,
            prompt_tokens=1,
            priority=RequestPriority.CRITICAL,
            arrival_time=now,
        ),
        RequestMetadata(
            request_id="req-batch-low",
            trace_id="trace-batch-low",
            model_id="Qwen2-7B",
            prompt="low",
            max_tokens=8,
            prompt_tokens=1,
            priority=RequestPriority.LOW,
            arrival_time=now,
        ),
    ]

    builder = IRBuilder()
    optimizer = IROptimizer()
    executor = DefaultIRExecutor(control_plane=manager)

    ir = builder.build_from_batch(requests)
    optimized = optimizer.optimize(ir)
    results = await executor.execute(optimized)

    assert "cmd-req-batch-critical" in results
    assert "cmd-req-batch-low" in results
    assert results["cmd-req-batch-critical"]["output_text"] == "cp-http-req-batch-critical"
    assert results["cmd-req-batch-low"]["output_text"] == "cp-http-req-batch-low"
