from __future__ import annotations

from django.db import models

from control_plane.common import TimestampedModel


class TrainingProjectStatus(models.TextChoices):
    IMPORTING = "IMPORTING"
    ARCHAEOLOGY = "ARCHAEOLOGY"
    NEEDS_REPAIR = "NEEDS_REPAIR"
    BASELINING = "BASELINING"
    READY = "READY"
    OVERNIGHT_RUNNING = "OVERNIGHT_RUNNING"
    PAUSED = "PAUSED"
    FAILED = "FAILED"
    FINISHED = "FINISHED"
    EVOLVING = "EVOLVING"


class DatasetValidationStatus(models.TextChoices):
    VALID = "VALID"
    WARNING = "WARNING"
    BLOCKED = "BLOCKED"
    UNKNOWN = "UNKNOWN"


class CheckpointType(models.TextChoices):
    BASE = "BASE"
    IMPORTED = "IMPORTED"
    CHAMPION = "CHAMPION"
    CHALLENGER = "CHALLENGER"
    INTERMEDIATE = "INTERMEDIATE"


class CheckpointValidityStatus(models.TextChoices):
    UNKNOWN = "UNKNOWN"
    VALID = "VALID"
    INVALID = "INVALID"
    PARTIAL = "PARTIAL"
    CORRUPT = "CORRUPT"


class ExperimentStatus(models.TextChoices):
    PROPOSED = "PROPOSED"
    QUEUED = "QUEUED"
    RUNNING = "RUNNING"
    EVALUATING = "EVALUATING"
    VALIDATING = "VALIDATING"
    PROMOTED = "PROMOTED"
    REJECTED = "REJECTED"
    INCONCLUSIVE = "INCONCLUSIVE"
    FAILED = "FAILED"
    CANCELLED = "CANCELLED"
    SUPERSEDED = "SUPERSEDED"


class TrainingRunStatus(models.TextChoices):
    QUEUED = "QUEUED"
    STARTING = "STARTING"
    RUNNING = "RUNNING"
    CHECKPOINTING = "CHECKPOINTING"
    SUCCEEDED = "SUCCEEDED"
    FAILED = "FAILED"
    OOM = "OOM"
    TIMEOUT = "TIMEOUT"
    CANCELLED = "CANCELLED"
    INTERRUPTED = "INTERRUPTED"


class EvaluationGroup(models.TextChoices):
    PRIMARY = "PRIMARY"
    CAPABILITY_SUBSET = "CAPABILITY_SUBSET"
    HOLDOUT = "HOLDOUT"
    REGRESSION = "REGRESSION"
    ADVERSARIAL = "ADVERSARIAL"
    FORMAT = "FORMAT"
    PERFORMANCE = "PERFORMANCE"


class EvaluationRunStatus(models.TextChoices):
    QUEUED = "QUEUED"
    RUNNING = "RUNNING"
    SUCCEEDED = "SUCCEEDED"
    FAILED = "FAILED"
    INVALID = "INVALID"


class ProgramStatus(models.TextChoices):
    CREATED = "CREATED"
    ARCHAEOLOGY = "ARCHAEOLOGY"
    BASELINING = "BASELINING"
    PLANNING = "PLANNING"
    RUNNING_EXPERIMENT = "RUNNING_EXPERIMENT"
    EVALUATING = "EVALUATING"
    ADAPTING = "ADAPTING"
    FINALIZING = "FINALIZING"
    COMPLETED = "COMPLETED"
    PARTIAL = "PARTIAL"
    FAILED = "FAILED"
    PAUSED = "PAUSED"


class PromotionDecision(models.TextChoices):
    PROMOTE = "PROMOTE"
    REJECT = "REJECT"
    INCONCLUSIVE = "INCONCLUSIVE"
    REQUIRE_REPLICATION = "REQUIRE_REPLICATION"


class Conclusion(models.TextChoices):
    SUPPORTED = "SUPPORTED"
    WEAKLY_SUPPORTED = "WEAKLY_SUPPORTED"
    REFUTED = "REFUTED"
    AMBIGUOUS = "AMBIGUOUS"
    EXECUTION_FAILED = "EXECUTION_FAILED"


class TrainingProject(TimestampedModel):
    project = models.ForeignKey("projects.Project", on_delete=models.PROTECT, related_name="training_projects")
    studio_type = models.CharField(max_length=32, default="MODEL")
    name = models.CharField(max_length=200)
    slug = models.SlugField(max_length=120, unique=True)
    description = models.TextField(blank=True)
    goal = models.TextField()
    capability_target = models.TextField(blank=True)
    model_family = models.CharField(max_length=160, blank=True)
    model_size = models.CharField(max_length=80, blank=True)
    base_model = models.TextField(blank=True)
    repository_path = models.TextField(blank=True)
    repository_reference = models.TextField(blank=True)
    working_directory = models.TextField(blank=True)
    status = models.CharField(max_length=32, choices=TrainingProjectStatus.choices, default=TrainingProjectStatus.IMPORTING)
    default_profile = models.CharField(max_length=120, default="guard")
    training_backend = models.CharField(max_length=120, default="guard_subprocess")
    current_champion = models.ForeignKey("ModelCheckpoint", on_delete=models.SET_NULL, null=True, blank=True, related_name="champion_for_projects")
    baseline_evaluation = models.ForeignKey("EvaluationRun", on_delete=models.SET_NULL, null=True, blank=True, related_name="baseline_for_projects")
    metadata = models.JSONField(default=dict, blank=True)


class Dataset(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="datasets")
    name = models.CharField(max_length=200)
    description = models.TextField(blank=True)
    metadata = models.JSONField(default=dict, blank=True)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["training_project", "name"], name="unique_model_studio_dataset")]


class DatasetVersion(TimestampedModel):
    dataset = models.ForeignKey(Dataset, on_delete=models.CASCADE, related_name="versions")
    version = models.CharField(max_length=120)
    parent_version = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="children")
    manifest_reference = models.TextField()
    content_hash = models.CharField(max_length=128)
    record_count = models.PositiveIntegerField(null=True, blank=True)
    split_metadata = models.JSONField(default=dict, blank=True)
    source_metadata = models.JSONField(default=dict, blank=True)
    generation_metadata = models.JSONField(default=dict, blank=True)
    tags = models.JSONField(default=list, blank=True)
    validation_status = models.CharField(max_length=16, choices=DatasetValidationStatus.choices, default=DatasetValidationStatus.UNKNOWN)
    contamination_status = models.CharField(max_length=16, choices=DatasetValidationStatus.choices, default=DatasetValidationStatus.UNKNOWN)
    immutable = models.BooleanField(default=False)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["dataset", "version"], name="unique_model_studio_dataset_version")]

    def save(self, *args, **kwargs):
        if self.pk and self.immutable:
            original = type(self).objects.get(pk=self.pk)
            immutable_fields = ["version", "parent_version_id", "manifest_reference", "content_hash", "record_count", "split_metadata", "source_metadata", "generation_metadata", "tags"]
            if any(getattr(self, field) != getattr(original, field) for field in immutable_fields):
                raise ValueError("DatasetVersion is immutable after completed training use.")
        super().save(*args, **kwargs)


class TrainingRecipe(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="recipes")
    name = models.CharField(max_length=200)
    configuration = models.JSONField(default=dict)
    recipe_hash = models.CharField(max_length=128, unique=True)
    immutable = models.BooleanField(default=False)

    def save(self, *args, **kwargs):
        if self.pk and self.immutable:
            original = type(self).objects.get(pk=self.pk)
            if self.configuration != original.configuration or self.recipe_hash != original.recipe_hash:
                raise ValueError("TrainingRecipe is immutable after a TrainingRun starts.")
        super().save(*args, **kwargs)


class TrainingExperiment(TimestampedModel):
    experiment_id = models.CharField(max_length=120, unique=True)
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="experiments")
    title = models.CharField(max_length=255)
    hypothesis = models.TextField()
    reasoning = models.TextField()
    intervention = models.JSONField(default=dict)
    controls = models.JSONField(default=dict)
    expected_result = models.TextField()
    primary_success_metric = models.CharField(max_length=160)
    success_threshold = models.JSONField(default=dict)
    regression_constraints = models.JSONField(default=dict)
    rejection_condition = models.TextField()
    ambiguity_policy = models.TextField()
    estimated_runtime_seconds = models.PositiveIntegerField(default=0)
    maximum_runtime_seconds = models.PositiveIntegerField(default=0)
    compute_budget = models.JSONField(default=dict)
    parent_experiment = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="derived_experiments")
    derived_from_failure_cluster = models.ForeignKey("FailureCluster", on_delete=models.SET_NULL, null=True, blank=True, related_name="experiments")
    status = models.CharField(max_length=32, choices=ExperimentStatus.choices, default=ExperimentStatus.PROPOSED)
    priority = models.FloatField(default=0)
    expected_information_gain = models.FloatField(default=0)
    expected_improvement = models.FloatField(default=0)
    estimated_compute_cost = models.FloatField(default=0)
    experiment_value_score = models.FloatField(default=0)
    fingerprint = models.CharField(max_length=128)
    created_by_agent = models.CharField(max_length=120, default="MODEL_DIRECTOR")
    approved_by_model_director = models.BooleanField(default=False)
    result_summary = models.TextField(blank=True)
    conclusion = models.CharField(max_length=32, choices=Conclusion.choices, blank=True)

    class Meta:
        indexes = [models.Index(fields=["training_project", "fingerprint"])]


class ExperimentDependency(TimestampedModel):
    experiment = models.ForeignKey(TrainingExperiment, on_delete=models.CASCADE, related_name="dependencies")
    depends_on = models.ForeignKey(TrainingExperiment, on_delete=models.CASCADE, related_name="dependents")
    required_conclusions = models.JSONField(default=list, blank=True)
    rationale = models.TextField(blank=True)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["experiment", "depends_on"], name="unique_experiment_dependency")]


class ModelCheckpoint(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="checkpoints")
    name = models.CharField(max_length=255)
    checkpoint_type = models.CharField(max_length=32, choices=CheckpointType.choices)
    reference = models.TextField()
    content_hash = models.CharField(max_length=128, blank=True)
    base_checkpoint = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="derived_checkpoints")
    training_run = models.ForeignKey("TrainingRun", on_delete=models.SET_NULL, null=True, blank=True, related_name="output_checkpoints")
    recipe = models.ForeignKey(TrainingRecipe, on_delete=models.SET_NULL, null=True, blank=True, related_name="checkpoints")
    dataset_versions = models.ManyToManyField(DatasetVersion, blank=True, related_name="checkpoints")
    parameter_count = models.BigIntegerField(null=True, blank=True)
    dtype = models.CharField(max_length=80, blank=True)
    adapter_type = models.CharField(max_length=80, blank=True)
    quantization = models.CharField(max_length=80, blank=True)
    validity_status = models.CharField(max_length=16, choices=CheckpointValidityStatus.choices, default=CheckpointValidityStatus.UNKNOWN)
    load_verified = models.BooleanField(default=False)
    evaluation_status = models.CharField(max_length=32, default="NOT_EVALUATED")
    metadata = models.JSONField(default=dict, blank=True)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["training_project", "reference"], name="unique_model_checkpoint_reference")]


class TrainingRun(TimestampedModel):
    experiment = models.ForeignKey(TrainingExperiment, on_delete=models.PROTECT, related_name="runs")
    recipe = models.ForeignKey(TrainingRecipe, on_delete=models.PROTECT, related_name="runs")
    input_checkpoint = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="input_runs")
    output_checkpoint = models.ForeignKey(ModelCheckpoint, on_delete=models.SET_NULL, null=True, blank=True, related_name="producing_run")
    status = models.CharField(max_length=32, choices=TrainingRunStatus.choices, default=TrainingRunStatus.QUEUED)
    command = models.JSONField(default=list, blank=True)
    working_directory = models.TextField(blank=True)
    environment_snapshot = models.JSONField(default=dict, blank=True)
    host = models.CharField(max_length=255, blank=True)
    gpu_device = models.CharField(max_length=255, blank=True)
    allocated_memory_mb = models.PositiveIntegerField(null=True, blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    wall_seconds = models.FloatField(default=0)
    exit_code = models.IntegerField(null=True, blank=True)
    stdout_reference = models.TextField(blank=True)
    stderr_reference = models.TextField(blank=True)
    training_log_reference = models.TextField(blank=True)
    peak_memory_mb = models.PositiveIntegerField(null=True, blank=True)
    gpu_utilization = models.FloatField(null=True, blank=True)
    failure_category = models.CharField(max_length=80, blank=True)
    failure_details = models.TextField(blank=True)
    resume_source = models.TextField(blank=True)
    retry_count = models.PositiveIntegerField(default=0)
    pid = models.IntegerField(null=True, blank=True)
    heartbeat_at = models.DateTimeField(null=True, blank=True)


class EvaluationSuite(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="evaluation_suites")
    name = models.CharField(max_length=200)
    description = models.TextField(blank=True)
    metadata = models.JSONField(default=dict, blank=True)


class EvaluationSuiteVersion(TimestampedModel):
    suite = models.ForeignKey(EvaluationSuite, on_delete=models.CASCADE, related_name="versions")
    version = models.CharField(max_length=120)
    reference = models.TextField()
    content_hash = models.CharField(max_length=128)
    command_template = models.JSONField(default=list, blank=True)
    groups = models.JSONField(default=list, blank=True)
    immutable = models.BooleanField(default=False)
    integrity_status = models.CharField(max_length=16, choices=DatasetValidationStatus.choices, default=DatasetValidationStatus.UNKNOWN)
    integrity_evidence = models.JSONField(default=dict, blank=True)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["suite", "version"], name="unique_evaluation_suite_version")]


class EvaluationRun(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="evaluation_runs")
    checkpoint = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="evaluation_runs")
    suite_version = models.ForeignKey(EvaluationSuiteVersion, on_delete=models.PROTECT, related_name="runs")
    status = models.CharField(max_length=16, choices=EvaluationRunStatus.choices, default=EvaluationRunStatus.QUEUED)
    command = models.JSONField(default=list, blank=True)
    output_reference = models.TextField(blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    wall_seconds = models.FloatField(default=0)
    parser_version = models.CharField(max_length=120, blank=True)
    integrity_evidence = models.JSONField(default=dict, blank=True)
    summary = models.JSONField(default=dict, blank=True)
    failure_details = models.TextField(blank=True)


class BenchmarkResult(TimestampedModel):
    evaluation_run = models.ForeignKey(EvaluationRun, on_delete=models.CASCADE, related_name="results")
    group = models.CharField(max_length=32, choices=EvaluationGroup.choices, default=EvaluationGroup.PRIMARY)
    metric = models.CharField(max_length=160)
    value = models.FloatField()
    unit = models.CharField(max_length=80, blank=True)
    subset = models.CharField(max_length=160, blank=True)
    sample_count = models.PositiveIntegerField(null=True, blank=True)
    passed = models.BooleanField(null=True, blank=True)
    provenance = models.JSONField(default=dict, blank=True)


class RegressionBankItem(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="regression_bank_items")
    reference = models.TextField()
    content_hash = models.CharField(max_length=128, blank=True)
    category = models.CharField(max_length=160, blank=True)
    kind = models.CharField(max_length=80, default="IMPORTED")
    provenance = models.JSONField(default=dict, blank=True)
    evaluation_only = models.BooleanField(default=True)


class FailureCluster(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="failure_clusters")
    name = models.CharField(max_length=200)
    description = models.TextField(blank=True)
    failure_type = models.CharField(max_length=160, blank=True)
    severity = models.CharField(max_length=32, default="UNKNOWN")
    sample_count = models.PositiveIntegerField(default=0)
    representative_examples = models.JSONField(default=list, blank=True)
    affected_benchmarks = models.JSONField(default=list, blank=True)
    suspected_causes = models.JSONField(default=list, blank=True)
    confidence = models.FloatField(default=0)
    training_data_coverage = models.JSONField(default=dict, blank=True)
    priority = models.FloatField(default=0)


class ModelPromotionPolicy(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="promotion_policies")
    name = models.CharField(max_length=160)
    version = models.CharField(max_length=80)
    criteria = models.JSONField(default=dict)
    active = models.BooleanField(default=True)


class ModelPromotionDecision(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="promotion_decisions")
    from_champion = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="promotion_sources")
    candidate = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="promotion_candidates")
    experiment = models.ForeignKey(TrainingExperiment, on_delete=models.PROTECT, related_name="promotion_decisions")
    decision = models.CharField(max_length=32, choices=PromotionDecision.choices)
    policy = models.ForeignKey(ModelPromotionPolicy, on_delete=models.PROTECT, related_name="decisions")
    evaluation_evidence = models.JSONField(default=dict)
    judge_result = models.JSONField(default=dict)
    reason = models.TextField()
    judge_actor = models.CharField(max_length=120, default="MODEL_JUDGE")


class OvernightTrainingProgram(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="overnight_programs")
    starting_champion = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="starting_programs")
    ending_champion = models.ForeignKey(ModelCheckpoint, on_delete=models.SET_NULL, null=True, blank=True, related_name="ending_programs")
    status = models.CharField(max_length=32, choices=ProgramStatus.choices, default=ProgramStatus.CREATED)
    start_time = models.DateTimeField(null=True, blank=True)
    deadline = models.DateTimeField()
    maximum_wall_seconds = models.PositiveIntegerField()
    maximum_training_runs = models.PositiveIntegerField(default=8)
    maximum_failed_runs = models.PositiveIntegerField(default=3)
    maximum_single_run_seconds = models.PositiveIntegerField()
    evaluation_reserve_seconds = models.PositiveIntegerField()
    allowed_experiment_types = models.JSONField(default=list, blank=True)
    promotion_policy = models.ForeignKey(ModelPromotionPolicy, on_delete=models.PROTECT, related_name="programs")
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    termination_reason = models.TextField(blank=True)
    pause_after_current_run = models.BooleanField(default=False)
    telemetry = models.JSONField(default=dict, blank=True)


class OvernightResearchReport(TimestampedModel):
    program = models.OneToOneField(OvernightTrainingProgram, on_delete=models.CASCADE, related_name="report")
    markdown = models.TextField()
    payload = models.JSONField(default=dict)
    artifact_reference = models.TextField(blank=True)


class ModelStudioArtifact(TimestampedModel):
    training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="artifacts")
    artifact_type = models.CharField(max_length=120)
    name = models.CharField(max_length=255)
    content = models.JSONField(default=dict)
    readable = models.TextField(blank=True)
    source_reference = models.TextField(blank=True)