import json

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

from control_plane.model_studio.models import DatasetVersion, TrainingProject
from control_plane.model_studio.services import ModelStudioService
from model_router.providers import providers_from_resources
from model_router.router import ModelRouter


class Command(BaseCommand):
    help = "Ask the Qwen Dataset Curator for a structured, non-mutating Guard dataset improvement proposal."

    def add_arguments(self, parser):
        parser.add_argument("--project", required=True, help="TrainingProject slug")
        parser.add_argument("--status", default="WARNING", help="Dataset validation status to review")

    def handle(self, *args, **options):
        project = TrainingProject.objects.filter(slug=options["project"]).first()
        if project is None:
            raise CommandError("TrainingProject not found.")
        versions = list(DatasetVersion.objects.filter(dataset__training_project=project, validation_status=options["status"]).order_by("-record_count"))
        if not versions:
            raise CommandError("No matching DatasetVersions to curate.")
        router = ModelRouter(providers_from_resources(), persist_requests=True)
        proposal = ModelStudioService(router=router).propose_dataset_curation(project, versions)
        self.stdout.write(json.dumps({"proposal_id": str(proposal.id), "title": proposal.title, "operations": proposal.proposed_operations, "validation_plan": proposal.validation_plan}, indent=2, default=str))