from pydantic import BaseModel, Field
from typing import Dict, Any, List, Optional, Union


class ClaudeJobRequest(BaseModel):
    """Request model for Claude to start or manage a job"""
    job_id: str = Field(..., description="The ID of the job to manage")
    action: str = Field(..., description="Action to perform: start, stop, restart, status")
    namespace: Optional[str] = Field("development", description="Nomad namespace")
    purge: Optional[bool] = Field(False, description="Whether to purge the job when stopping")


class ClaudeJobSpecification(BaseModel):
    """Simplified job specification for Claude to create a new job"""
    job_id: str = Field(..., description="The ID for the new job")
    name: Optional[str] = Field(None, description="Name of the job (defaults to job_id)")
    type: str = Field("service", description="Job type: service, batch, or system")
    datacenters: List[str] = Field(["jm"], description="List of datacenters")
    namespace: str = Field("development", description="Nomad namespace")
    docker_image: str = Field(..., description="Docker image to run")
    count: int = Field(1, description="Number of instances to run")
    cpu: int = Field(100, description="CPU resources in MHz")
    memory: int = Field(128, description="Memory in MB")
    ports: Optional[List[Dict[str, Any]]] = Field(None, description="Port mappings")
    env_vars: Optional[Dict[str, str]] = Field(None, description="Environment variables")
    
    def to_nomad_job_spec(self) -> Dict[str, Any]:
        """Convert to Nomad job specification format"""
        # Create a task with the specified Docker image
        task = {
            "Name": "app",
            "Driver": "docker",
            "Config": {
                "image": self.docker_image,
            },
            "Resources": {
                "CPU": self.cpu,
                "MemoryMB": self.memory
            }
        }
        
        # Add environment variables if specified
        if self.env_vars:
            task["Env"] = self.env_vars
            
        # Create network configuration
        network = {}
        if self.ports:
            network["DynamicPorts"] = self.ports
            task["Config"]["ports"] = [port["Label"] for port in self.ports]
            
        # Create the full job specification
        job_spec = {
            "ID": self.job_id,
            "Name": self.name or self.job_id,
            "Type": self.type,
            "Datacenters": self.datacenters,
            "Namespace": self.namespace,
            "TaskGroups": [
                {
                    "Name": "app",
                    "Count": self.count,
                    "Tasks": [task],
                    "Networks": [network] if network else []
                }
            ]
        }
        
        return job_spec


class ClaudeJobResponse(BaseModel):
    """Response model for Claude job operations"""
    success: bool = Field(..., description="Whether the operation was successful")
    job_id: str = Field(..., description="The ID of the job")
    status: str = Field(..., description="Current status of the job")
    message: str = Field(..., description="Human-readable message about the operation")
    details: Optional[Dict[str, Any]] = Field(None, description="Additional details about the job")