import json

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

from control_plane.model_studio.models import DatasetCurationProposal
from control_plane.model_studio.services import ModelStudioService


class Command(BaseCommand):
    help = "Materialize a Qwen curation proposal into a new immutable Spark dataset version and source-hash splits."

    def add_arguments(self, parser):
        parser.add_argument("--proposal", required=True)
        parser.add_argument("--output-directory", required=True)
        parser.add_argument("--ssh-alias", default="spark")
        parser.add_argument("--strict-schema-repair", action="store_true")

    def handle(self, *args, **options):
        proposal = DatasetCurationProposal.objects.filter(id=options["proposal"]).first()
        if proposal is None:
            raise CommandError("DatasetCurationProposal not found.")
        training_project = proposal.source_versions.first().dataset.training_project if proposal.source_versions.exists() else None
        if training_project is None:
            raise CommandError("Proposal has no source versions.")
        version = ModelStudioService().materialize_spark_curation(training_project, proposal, output_directory=options["output_directory"], ssh_alias=options["ssh_alias"], strict=options["strict_schema_repair"])
        self.stdout.write(json.dumps({"dataset_version": str(version.id), "reference": version.manifest_reference, "records": version.record_count, "splits": version.split_metadata}, indent=2))