from __future__ import annotations

import json
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

from django.core.management.base import BaseCommand, CommandError

from control_plane.resources.models import Resource
from model_router.providers import QwenProvider
from model_router.router import ModelCapability, ModelRequestContract


class Command(BaseCommand):
    help = "Probe Qwen/local_inference health and bounded concurrent completion capacity."

    def add_arguments(self, parser):
        parser.add_argument("--requests", type=int, default=1)
        parser.add_argument("--concurrency", type=int, default=1)
        parser.add_argument("--prompt", default="Return only valid JSON: {\"ok\": true, \"model\": \"qwen\"}.")
        parser.add_argument("--token-budget", type=int, default=256)

    def handle(self, *args, **options):
        resource = Resource.objects.filter(provider="local_inference", is_active=True).first()
        if resource is None:
            raise CommandError("No Qwen/local_inference resource configured. Run seed_spark_resources first.")
        provider = QwenProvider(resource)
        before_health = provider.health()
        total = max(1, int(options["requests"]))
        concurrency = max(1, int(options["concurrency"]))
        started = time.monotonic()

        def call(index: int) -> dict[str, object]:
            call_started = time.monotonic()
            try:
                response = provider.complete(
                    ModelRequestContract(
                        purpose=ModelCapability.REASONING,
                        prompt=f"{options['prompt']} Request index: {index}",
                        token_budget=int(options["token_budget"]),
                    )
                )
                return {
                    "index": index,
                    "status": "COMPLETE",
                    "latency_ms": int((time.monotonic() - call_started) * 1000),
                    "content_chars": len(response.content),
                    "model": response.model,
                    "attempts": response.metadata.get("attempts"),
                }
            except Exception as exc:
                return {
                    "index": index,
                    "status": "FAILED",
                    "latency_ms": int((time.monotonic() - call_started) * 1000),
                    "failure": str(exc),
                }

        results = []
        with ThreadPoolExecutor(max_workers=concurrency) as executor:
            futures = {executor.submit(call, index): index for index in range(1, total + 1)}
            for future in as_completed(futures):
                results.append(future.result())
        results.sort(key=lambda row: row["index"])
        after_health = provider.health()
        completed = sum(1 for row in results if row["status"] == "COMPLETE")
        failed = total - completed
        latencies = [int(row["latency_ms"]) for row in results]
        output = {
            "before_health": before_health,
            "after_health": after_health,
            "requests": total,
            "concurrency": concurrency,
            "completed": completed,
            "failed": failed,
            "runtime_seconds": round(time.monotonic() - started, 2),
            "latency_ms": {"min": min(latencies), "max": max(latencies), "avg": round(sum(latencies) / len(latencies), 1)},
            "results": results,
        }
        self.stdout.write(json.dumps(output, indent=2))
        if failed:
            raise CommandError(f"Qwen probe failed {failed}/{total} requests")