"""Inference suite — single request + throughput measurement."""

from __future__ import annotations

import asyncio
import json
from typing import TYPE_CHECKING

from tests.remote.suites.base import BaseSuite, StepStatus, SuiteResult

if TYPE_CHECKING:
    from tests.remote.runner import BackgroundProcess, RemoteTestRunner

SERVER_PORT = 4000
COMPLETIONS_URL = f"http://localhost:{SERVER_PORT}/v1/chat/completions"

# Minimal chat payload
CHAT_PAYLOAD = json.dumps(
    {
        "model": "auto",
        "messages": [{"role": "user", "content": "Say hello in exactly 5 words."}],
        "max_tokens": 50,
    }
)


class InferenceSuite(BaseSuite):
    name = "inference"
    description = "Run inference requests and measure latency and throughput"

    async def execute(self, runner: RemoteTestRunner, result: SuiteResult) -> None:
        proc: BackgroundProcess | None = None

        try:
            # Step 1: Start server
            proc = await runner.run_background("llmhosts serve --no-tui")
            if not proc.pid:
                self._add_step(
                    result,
                    "start_server",
                    "Start server for inference",
                    status=StepStatus.FAILED,
                    error_message="Failed to start server",
                )
                return

            # Wait for server ready
            ready = False
            for _ in range(15):
                await asyncio.sleep(2)
                curl = self._health_cmd(runner)
                r = await runner.run_command(curl, timeout=10.0)
                if r.success and ("ok" in r.stdout.lower() or "healthy" in r.stdout.lower()):
                    ready = True
                    break

            if not ready:
                self._add_step(
                    result,
                    "start_server",
                    "Wait for server health",
                    status=StepStatus.FAILED,
                    error_message="Server did not become healthy in 30s",
                )
                return

            self._add_step(
                result,
                "start_server",
                f"Server started (PID={proc.pid})",
                status=StepStatus.PASSED,
            )

            # Step 2: Single inference request
            curl = self._chat_cmd(runner)
            r = await runner.run_command(curl, timeout=120.0)

            if r.success and r.stdout:
                try:
                    resp = json.loads(r.stdout)
                    content = resp.get("choices", [{}])[0].get("message", {}).get("content", "")
                    tokens = resp.get("usage", {})
                    self._add_step(
                        result,
                        "single_request",
                        f"Single inference: {r.duration_s}s",
                        status=StepStatus.PASSED,
                        stdout=f"Response: {content}\nTokens: {json.dumps(tokens)}",
                        duration_s=r.duration_s,
                    )
                except json.JSONDecodeError:
                    self._add_step(
                        result,
                        "single_request",
                        "Single inference request",
                        status=StepStatus.FAILED,
                        stdout=r.stdout[:500],
                        error_message="Response is not valid JSON",
                        duration_s=r.duration_s,
                    )
                    return
            else:
                self._add_step(
                    result,
                    "single_request",
                    "Single inference request",
                    status=StepStatus.FAILED,
                    stdout=r.stdout,
                    stderr=r.stderr,
                    duration_s=r.duration_s,
                    error_message=f"Request failed (exit={r.exit_code})",
                )
                return

            # Step 3: Throughput test — 5 sequential requests
            total_tokens = 0
            total_time = 0.0
            success_count = 0

            for _i in range(5):
                r = await runner.run_command(curl, timeout=120.0)
                if r.success and r.stdout:
                    try:
                        resp = json.loads(r.stdout)
                        usage = resp.get("usage", {})
                        completion_tokens = usage.get("completion_tokens", 0)
                        total_tokens += completion_tokens
                        total_time += r.duration_s
                        success_count += 1
                    except json.JSONDecodeError:
                        total_time += r.duration_s

            if success_count > 0 and total_time > 0:
                tps = total_tokens / total_time if total_tokens > 0 else 0
                self._add_step(
                    result,
                    "throughput",
                    f"Throughput: {success_count}/5 OK, {total_tokens} tokens in {total_time:.1f}s ({tps:.1f} tok/s)",
                    status=StepStatus.PASSED,
                    stdout=f"Requests: {success_count}/5\nTotal tokens: {total_tokens}\nTotal time: {total_time:.1f}s\nTokens/sec: {tps:.1f}",
                    duration_s=total_time,
                )
            else:
                self._add_step(
                    result,
                    "throughput",
                    "Throughput test",
                    status=StepStatus.FAILED,
                    error_message=f"Only {success_count}/5 requests succeeded",
                    duration_s=total_time,
                )

        finally:
            if proc is not None:
                stop_result = await runner.stop_background(proc)
                self._add_step(
                    result,
                    "stop_server",
                    "Stop server",
                    status=StepStatus.PASSED if stop_result.exit_code >= 0 else StepStatus.FAILED,
                    duration_s=stop_result.duration_s,
                )

    @staticmethod
    def _health_cmd(runner: RemoteTestRunner) -> str:
        url = f"http://localhost:{SERVER_PORT}/health"
        if runner.device.shell == "powershell":
            return f'(Invoke-WebRequest -Uri "{url}" -UseBasicParsing -TimeoutSec 5).Content'
        return f"curl -sf --connect-timeout 5 {url}"

    @staticmethod
    def _chat_cmd(runner: RemoteTestRunner) -> str:
        """Build a chat completions curl command for the device shell."""
        if runner.device.shell == "powershell":
            # Use Invoke-RestMethod for PowerShell
            payload_escaped = CHAT_PAYLOAD.replace('"', '`"')
            return (
                f'Invoke-RestMethod -Uri "{COMPLETIONS_URL}" '
                f'-Method POST -ContentType "application/json" '
                f'-Body "{payload_escaped}" -TimeoutSec 120 | ConvertTo-Json -Depth 10'
            )
        return (
            f"curl -sf --connect-timeout 5 -m 120 "
            f"-H 'Content-Type: application/json' "
            f"-d '{CHAT_PAYLOAD}' "
            f"{COMPLETIONS_URL}"
        )
