from __future__ import annotations

from django.core.management.base import BaseCommand
from django.utils import timezone

from control_plane.resources.models import Resource
from model_router.providers import QwenProvider, SolProvider
from research.searxng import SearxngSearchClient


class Command(BaseCommand):
    help = "Check configured model provider health without crashing the control plane."

    def handle(self, *args, **options):
        resources = sorted(Resource.objects.filter(is_active=True, provider__in=["local_inference", "opencode", "searxng"]), key=self.sort_key)
        for resource in resources:
            if resource.provider == "opencode":
                status = SolProvider(resource).health()
            elif resource.provider == "local_inference":
                status = QwenProvider(resource).health()
            elif resource.provider == "searxng":
                status = SearxngSearchClient(endpoint_url=str(resource.config.get("endpoint_url", "http://127.0.0.1:8080"))).health()
            else:
                status = "UNAVAILABLE"
            resource.health_status = status
            resource.last_health_check_at = timezone.now()
            resource.save(update_fields=["health_status", "last_health_check_at", "updated_at"])
            self.stdout.write(f"{resource.name}: {status}")

    def sort_key(self, resource: Resource) -> tuple[int, str]:
        model_key = str(resource.config.get("model_key") or "").lower()
        if resource.provider == "local_inference":
            model_key = "qwen"
        if resource.provider == "searxng":
            model_key = "searxng"
        order = {"qwen": 0, "sol": 1, "terra": 2, "luna": 3, "searxng": 4}
        return order.get(model_key, 100), resource.name