"""
FireVM Python Client
"""

import requests
import time
from typing import Optional, Dict, List
from dataclasses import dataclass


class VMQuotaExceeded(Exception):
    """Raised when VM quota is exceeded"""
    def __init__(self, message, current_vms, max_vms, tier):
        self.message = message
        self.current_vms = current_vms
        self.max_vms = max_vms
        self.tier = tier
        super().__init__(self.message)
    
    def __str__(self):
        return f"{self.message} (Current: {self.current_vms}/{self.max_vms} on '{self.tier}' tier)"


@dataclass
class ExecutionResult:
    """Result of code execution"""
    stdout: str
    stderr: str
    exit_code: int
    execution_time: float
    error: Optional[str] = None


@dataclass
class VM:
    """VM instance"""
    vm_id: str
    name: str
    state: str
    ip_address: Optional[str]
    created_at: str
    
    def __str__(self):
        return f"VM({self.vm_id}, state={self.state})"


class FireVMClient:
    """
    Official FireVM Python SDK
    
    Example:
        >>> from firevm import FireVMClient
        >>> 
        >>> # Initialize client
        >>> client = FireVMClient(api_key="your_api_key")
        >>> 
        >>> # Create a VM
        >>> vm = client.create_vm(name="my-vm")
        >>> print(f"Created {vm}")
        >>> 
        >>> # Execute code
        >>> result = client.execute(vm.vm_id, "print('Hello World!')")
        >>> print(result.stdout)  # Hello World!
    """
    
    def __init__(self, api_key: str, base_url: str = "https://firevm-backend-345160625255.asia-northeast1.run.app"):
        """
        Initialize FireVM client
        
        Args:
            api_key: Your FireVM API key
            base_url: API base URL (default: production)
        """
        self.api_key = api_key
        self.base_url = base_url.rstrip('/')
        self.headers = {"Authorization": f"Bearer {api_key}"}
    
    def create_vm(self, name: str, tier: str = "small", wait: bool = True) -> VM:
        """
        Create a new VM
        
        Args:
            name: VM name
            tier: VM tier (small, medium, large)
            wait: Wait for VM to be ready
        
        Returns:
            VM instance
        
        Raises:
            VMQuotaExceeded: If user has reached their VM limit
            requests.HTTPError: For other API errors
        
        Example:
            >>> try:
            ...     vm = client.create_vm("my-vm")
            ... except VMQuotaExceeded as e:
            ...     print(f"Limit reached: {e}")
            ...     print(f"You have {e.current_vms}/{e.max_vms} VMs")
        """
        url = f"{self.base_url}/api/v1/vms/"
        data = {"name": name, "tier": tier, "image": "ubuntu-22.04"}
        
        response = requests.post(url, headers=self.headers, json=data)
        
        # Check for quota exceeded error (403)
        if response.status_code == 403:
            try:
                error_detail = response.json().get("detail", {})
                if isinstance(error_detail, dict) and error_detail.get("error") == "vm_limit_reached":
                    raise VMQuotaExceeded(
                        message=error_detail.get("message", "VM limit reached"),
                        current_vms=error_detail.get("current_vms", 0),
                        max_vms=error_detail.get("max_vms", 0),
                        tier=error_detail.get("tier", "unknown")
                    )
            except (ValueError, KeyError, AttributeError):
                pass  # Fall through to generic error handling
        
        # Raise for any HTTP errors
        response.raise_for_status()
        vm_data = response.json()
        
        vm = VM(
            vm_id=vm_data["vm_id"],
            name=vm_data["name"],
            state=vm_data["state"],
            ip_address=vm_data.get("infrastructure", {}).get("ip_address"),
            created_at=vm_data["created_at"]
        )
        
        if wait:
            vm = self._wait_for_vm(vm.vm_id)
        
        return vm
    
    def execute(self, vm_id: str, code: str, language: str = "python", 
                timeout: int = 30, wait: bool = True) -> ExecutionResult:
        """
        Execute code in a VM
        
        Args:
            vm_id: VM ID
            code: Code to execute
            language: Programming language (python, bash, nodejs)
            timeout: Execution timeout in seconds
            wait: Wait for execution to complete
        
        Returns:
            ExecutionResult with stdout, stderr, exit_code
        
        Example:
            >>> result = client.execute(vm.vm_id, "print(2 + 2)")
            >>> print(result.stdout)  # 4
            >>> print(result.exit_code)  # 0
        """
        # Submit execution job
        url = f"{self.base_url}/api/v1/vms/{vm_id}/execute"
        params = {"code": code, "language": language, "timeout": timeout}
        
        start_time = time.time()
        response = requests.post(url, headers=self.headers, params=params)
        response.raise_for_status()
        
        job_data = response.json()
        job_id = job_data["job_id"]
        
        if not wait:
            return ExecutionResult(
                stdout="",
                stderr="",
                exit_code=-1,
                execution_time=0,
                error="Job submitted but not waited"
            )
        
        # Poll for result
        job = self._wait_for_job(job_id, timeout + 10)
        execution_time = time.time() - start_time
        
        return ExecutionResult(
            stdout=job.get("stdout", ""),
            stderr=job.get("stderr", ""),
            exit_code=job.get("exit_code", -1),
            execution_time=execution_time,
            error=job.get("error")
        )
    
    def execute_python(self, vm_id: str, code: str, **kwargs) -> ExecutionResult:
        """Execute Python code (shorthand)"""
        return self.execute(vm_id, code, language="python", **kwargs)
    
    def execute_bash(self, vm_id: str, code: str, **kwargs) -> ExecutionResult:
        """Execute Bash code (shorthand)"""
        return self.execute(vm_id, code, language="bash", **kwargs)
    
    def _wait_for_vm(self, vm_id: str, timeout: int = 90) -> VM:
        """Wait for VM to be ready"""
        start = time.time()
        while time.time() - start < timeout:
            time.sleep(2)
            
            # Check MongoDB directly via execution to see if running
            # (workaround for get_vm validation issue)
            try:
                test_result = self.execute(
                    vm_id, 
                    "echo ready", 
                    language="bash",
                    timeout=5,
                    wait=True
                )
                if test_result.exit_code == 0:
                    return VM(
                        vm_id=vm_id,
                        name="",
                        state="running",
                        ip_address=None,
                        created_at=""
                    )
            except:
                continue
        
        # If we can't verify via execution, wait standard time
        time.sleep(max(0, 60 - (time.time() - start)))
        return VM(vm_id=vm_id, name="", state="running", ip_address=None, created_at="")
    
    def _wait_for_job(self, job_id: str, timeout: int = 60) -> Dict:
        """Wait for job to complete"""
        url = f"{self.base_url}/api/v1/jobs/{job_id}"
        start = time.time()
        
        while time.time() - start < timeout:
            response = requests.get(url, headers=self.headers)
            response.raise_for_status()
            
            job = response.json()
            state = job.get("state")
            
            if state in ["completed", "failed"]:
                return job
            
            time.sleep(1)
        
        raise TimeoutError(f"Job {job_id} did not complete within {timeout}s")
