from __future__ import annotations

import uuid

from django.db import models

from control_plane.common import TimestampedModel


class ResourceKind(models.TextChoices):
    MODEL = "MODEL"
    COMPUTE = "COMPUTE"
    MODEL_COMPUTE = "MODEL_COMPUTE"
    GIT = "GIT"
    ARTIFACT_STORAGE = "ARTIFACT_STORAGE"


class Resource(TimestampedModel):
    name = models.CharField(max_length=200, unique=True)
    kind = models.CharField(max_length=80, choices=ResourceKind.choices)
    compute = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="model_resources")
    provider = models.CharField(max_length=120, blank=True)
    roles = models.JSONField(default=list, blank=True)
    config = models.JSONField(default=dict, blank=True)
    health_status = models.CharField(max_length=32, default="UNKNOWN")
    last_health_check_at = models.DateTimeField(null=True, blank=True)
    is_active = models.BooleanField(default=True)


class ModelRequest(TimestampedModel):
    correlation_id = models.CharField(max_length=64, db_index=True, default=uuid.uuid4)
    project = models.ForeignKey("projects.Project", on_delete=models.CASCADE, null=True, blank=True)
    agent_version = models.ForeignKey("agents.AgentVersion", on_delete=models.SET_NULL, null=True, blank=True)
    logical_role = models.CharField(max_length=80, default="REASONING")
    model_resource = models.ForeignKey(Resource, on_delete=models.SET_NULL, null=True, blank=True, related_name="model_requests")
    provider = models.CharField(max_length=120, blank=True)
    model = models.CharField(max_length=120)
    priority = models.PositiveSmallIntegerField(default=50)
    token_budget = models.PositiveIntegerField(default=8000)
    status = models.CharField(max_length=32, default="QUEUED")
    request = models.JSONField(default=dict)
    response = models.JSONField(default=dict, blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    ended_at = models.DateTimeField(null=True, blank=True)
    latency_ms = models.PositiveIntegerField(null=True, blank=True)
    prompt_tokens = models.PositiveIntegerField(null=True, blank=True)
    completion_tokens = models.PositiveIntegerField(null=True, blank=True)
    failure_reason = models.TextField(blank=True)