# Copyright 2025 Foundry Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""FCP API client for spot instance management."""

import os
from typing import Dict, List, Optional
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from pathlib import Path

from .errors import APIError, ConfigError
from .catalog import resolver


class FCPClient:
    """Client for FCP spot instance operations."""
    
    def __init__(self, api_key: Optional[str] = None):
        # Try environment variables first
        self.api_key = api_key or os.getenv("FCP_API_KEY")
        self.project_id = os.getenv("FCP_PROJECT")
        
        # If not in env, try reading Flow config
        if not self.api_key or not self.project_id:
            config = self._read_flow_config()
            if config:
                self.api_key = self.api_key or config.get("api_key")
                self.project_id = self.project_id or config.get("project")
        
        # Validate we have required config
        if not self.api_key:
            raise ConfigError(
                "No API key found. Run 'flow init' or set FCP_API_KEY\n"
                "Note: If you've already run 'flow init', the API key may be in your system keychain.\n"
                "Set FCP_API_KEY environment variable to use it with mithril."
            )
            
        if not self.project_id:
            raise ConfigError("No project found. Run 'flow init' or set FCP_PROJECT")
            
        # Configure HTTP session with exponential backoff for transient errors
        self.session = requests.Session()
        retry = Retry(total=3, backoff_factor=0.3, status_forcelist=[502, 503, 504])
        self.session.mount("https://", HTTPAdapter(max_retries=retry))
        self.session.headers["Authorization"] = f"Bearer {self.api_key}"
        self.base_url = "https://api.mlfoundry.com/v2"
    
    def _read_flow_config(self) -> Optional[Dict]:
        """Read Flow SDK configuration if available."""
        config_path = Path.home() / ".flow" / "config.yaml"
        if not config_path.exists():
            return None
            
        try:
            import yaml
            with open(config_path) as f:
                return yaml.safe_load(f) or {}
        except Exception:
            # If YAML parsing fails, return None
            return None
    
    def _request(self, method: str, path: str, **kwargs) -> Dict:
        """Execute authenticated API request.
        
        Args:
            method: HTTP method
            path: API endpoint path
            **kwargs: Additional request parameters
            
        Returns:
            Parsed JSON response
            
        Raises:
            APIError: On HTTP 4xx/5xx responses
        """
        url = f"{self.base_url}{path}"
        response = self.session.request(method, url, **kwargs)
        
        if response.status_code >= 400:
            try:
                error = response.json()
                message = error.get("message", response.text)
            except:
                message = response.text
            raise APIError(response.status_code, message)
            
        return response.json()
    
    def _resolve_instance_type_name(self, instance_id: str) -> str:
        """Convert instance type ID to canonical name."""
        # Lazy-load instance type mappings
        if not hasattr(self, '_resolver_populated'):
            try:
                instance_types = self._request("GET", "/instance-types")
                resolver.populate_from_api(instance_types)
                self._resolver_populated = True
            except:
                # Fallback to static mappings on API failure
                self._resolver_populated = True
        
        return resolver.resolve(instance_id)
    
    def list_available(self) -> List[Dict]:
        """Query current spot market availability."""
        data = self._request("GET", "/spot/availability")
        # Handle both list and dict response formats
        return data if isinstance(data, list) else []
    
    def create_instance(self, 
                       instance_type: str,
                       ssh_key_path: Optional[str] = None,
                       region: Optional[str] = None) -> str:
        """Submit spot bid for specified instance type.
        
        Args:
            instance_type: Canonical instance type name
            ssh_key_path: Path to SSH public key file
            region: Target region constraint
            
        Returns:
            Bid ID for status tracking
            
        Raises:
            ConfigError: Missing SSH key
            APIError: No availability or bid rejection
        """
        # Load SSH public key
        ssh_key_path = ssh_key_path or os.getenv("FCP_SSH_KEY", os.path.expanduser("~/.ssh/id_rsa.pub"))
        try:
            with open(ssh_key_path, "r") as f:
                public_key = f.read().strip()
        except FileNotFoundError:
            raise ConfigError(f"SSH key not found: {ssh_key_path}")
        
        # Register SSH key with project
        ssh_key_id = self._ensure_ssh_key(public_key)
        
        # Match instance type to available auctions
        auctions = self.list_available()
        # Compare canonical names for type matching
        candidates = []
        for a in auctions:
            resolved_name = self._resolve_instance_type_name(a["instance_type"])
            if resolved_name == instance_type:
                candidates.append(a)
        
        if region:
            candidates = [a for a in candidates if a["region"] == region]
            
        if not candidates:
            raise APIError(404, f"No availability for {instance_type}")
            
        # Select lowest latency region by default
        auction = candidates[0]
        
        # Extract spot price from auction data
        price_str = auction.get("last_instance_price", auction.get("min_bid_price", "$1.00"))
        
        bid_data = {
            "project": self.project_id,
            "region": auction["region"],
            "name": f"mithril-{instance_type.replace('.', '-')}",
            "instance_quantity": 1,
            "limit_price": price_str,
            "auction_id": auction["fid"],  # Use auction ID, not instance type
            "ssh_keys": [ssh_key_id],
            "startup_script": "#!/bin/bash\necho 'Mithril instance ready'"
        }
        
        try:
            response = self._request("POST", "/spot/bids", json=bid_data)
            return response.get("fid", response.get("id", "unknown"))
        except APIError as e:
            raise APIError(e.status_code, f"Failed to create bid: {e.message}")
    
    def get_instance(self, instance_id: str) -> Dict:
        """Retrieve bid and instance status."""
        # Query bid status
        bid = self._request("GET", f"/spot/bids/{instance_id}")
        
        # Extract provisioned instance details if available
        instance = bid.get("instance")
        if not instance:
            return {
                "id": instance_id,
                "status": bid.get("status", "PENDING"),
                "type": bid.get("instance_type"),
                "region": bid.get("region")
            }
            
        return {
            "id": instance_id,
            "status": instance.get("status", "RUNNING"),
            "type": instance.get("instance_type"),
            "region": instance.get("region"),
            "ip": instance.get("public_ip"),
            "created": bid.get("created_at")
        }
    
    def delete_instance(self, instance_id: str) -> bool:
        """Cancel bid and terminate associated instance."""
        try:
            self._request("DELETE", f"/spot/bids/{instance_id}")
            return True
        except APIError as e:
            if e.status_code == 404:
                return False  # Idempotent deletion
            raise
    
    def list_my_instances(self) -> List[Dict]:
        """List active bids for project."""
        bids = self._request("GET", f"/spot/bids?project={self.project_id}")
        
        instances = []
        for bid in bids.get("bids", []):
            if bid.get("status") in ["PENDING", "RUNNING"]:
                instances.append({
                    "id": bid["id"],
                    "type": bid.get("instance_type"),
                    "status": bid.get("status"),
                    "region": bid.get("region")
                })
        return instances
    
    def _ensure_ssh_key(self, public_key: str) -> str:
        """Register SSH key if not already present.
        
        Args:
            public_key: OpenSSH format public key
            
        Returns:
            FCP key identifier
        """
        # Query project SSH keys
        keys = self._request("GET", f"/ssh-keys?project={self.project_id}")
        
        # Match by key fingerprint
        key_fingerprint = public_key.split()[1]  # Extract base64 key data
        for key in keys:
            if key_fingerprint in key.get("public_key", ""):
                return key["fid"]
        
        # Register new SSH key
        key_data = {
            "project": self.project_id,
            "name": "mithril-key",
            "public_key": public_key
        }
        response = self._request("POST", "/ssh-keys", json=key_data)
        return response["fid"]