import asyncio
import os
import json
import logging
import requests
from enum import Enum
from typing import TypedDict, Optional, List
from sagemaker_jupyterlab_extension_common.util.environment import EnvironmentDetector, Environment

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('SageMakerGenAIJupyterLabExtension')

class EnvironmentMetadata(TypedDict):
    """SageMaker environment metadata for origin validation"""
    domain_id: str
    project_id: str
    aws_region: str
    environment_id: str
    repository_name: str
    user_id: str
    dz_stage: str
    sm_domain_id: str
    sm_space_name: str
    sm_user_profile_name: str
    sm_project_path: str

class FilePaths(Enum):
    """File paths for various configurations"""
    BEARER_TOKEN = "~/.aws/sso/idc_access_token.json"
    Q_CUSTOMIZATION_ARN = "~/.aws/amazon_q/customization_arn.json"
    Q_DEV_PROFILE = "~/.aws/amazon_q/q_dev_profile.json"
    SETTINGS = "~/.aws/amazon_q/settings.json"
    CONFIG = "~/.aws/config"

class AuthMode(Enum):
    """Authentication modes for Amazon Q"""
    IDC = "IDC"  # Identity Center authentication
    IAM = "IAM"  # IAM authentication

def extract_bearer_token():
    """Extract bearer token from ~/.aws/sso/idc_access_token.json"""
    file_path = FilePaths.BEARER_TOKEN.value
    try:
        with open(os.path.expanduser(file_path), 'r') as f:
            content = json.load(f)
        
        token = content.get('idc_access_token')
        if not token or not token.strip():
            logger.error(f"No bearer token found in {file_path}")
            return None
        
        return token
    except FileNotFoundError:
        logger.error(f"Token file not found: {file_path}")
        return None
    except json.JSONDecodeError:
        logger.error(f"Invalid JSON in token file: {file_path}")
        return None
    except Exception as e:
        logger.error(f"Error reading token file: {e}")
        return None
    
def extract_q_customization_arn():
    """Extract customization ARN from ~/.aws/amazon_q/customization_arn.json"""
    file_path = os.path.expanduser(FilePaths.Q_CUSTOMIZATION_ARN.value)
    logger.info(f"Looking for customization file at: {file_path}")
    try:
        if not os.path.exists(file_path):
            logger.info(f"Customization file does not exist at: {file_path}")
            return None
  
        with open(file_path, "r") as f:
            content = json.load(f)
        customization_arn = content.get("customization_arn")
        if not customization_arn or not customization_arn.strip():
            logger.info(f"No customization ARN found in {file_path}")
            return None
        logger.info(f"Found Q customization ARN: {customization_arn}")
        return customization_arn
    except FileNotFoundError:
        logger.info(f"Customization file not found: {file_path}")
        return None
    except json.JSONDecodeError:
        logger.error(f"Invalid JSON in customization file: {file_path}")
        return None
    except Exception as e:
         logger.error(f"Error reading customization file: {e}")
         return None
    
def extract_q_settings():
    """Extract bearer token from ~/.aws/sso/idc_access_token.json"""
    file_path = FilePaths.Q_DEV_PROFILE.value
    try:
        with open(os.path.expanduser(file_path), 'r') as f:
            content = json.load(f)
        q_profile_arn = content.get('q_dev_profile_arn')
        if not q_profile_arn or not q_profile_arn.strip():
            logger.error(f"No Q profile ARN found in {file_path}")
            return None
        logger.info(f"Returning Q profile ARN: {q_profile_arn}")
        return q_profile_arn
    except FileNotFoundError:
        logger.error(f"Profile  not found: {file_path}")
        return None
    except json.JSONDecodeError:
        logger.error(f"Invalid JSON in token file: {file_path}")
        return None
    except Exception as e:
        logger.error(f"Error reading token file: {e}")
        return None

# Default to True, as mitigation to reading q settings file race condition
# This will be corrected when the q_enabled_file_watcher is started in WebSocketHandler 
def extract_q_enabled_status():
    """Extract q_enabled from ~/.aws/amazon_q/settings.json"""
    file_path = FilePaths.SETTINGS.value
    try:
        with open(os.path.expanduser(file_path), 'r') as f:
            content = json.load(f)
        q_enabled = content.get('q_enabled', True)
        logger.info(f"Returning q_enabled: {q_enabled}")
        return q_enabled
    except FileNotFoundError:
        logger.error(f"Settings file not found: {file_path}")
        return True
    except json.JSONDecodeError:
        logger.error(f"Invalid JSON in settings file: {file_path}")
        return True
    except Exception as e:
        logger.error(f"Error reading settings file: {e}")
        return True

def extract_auth_mode():
    """Extract auth mode from ~/.aws/amazon_q/settings.json"""
    file_path = FilePaths.SETTINGS.value
    try:
        with open(os.path.expanduser(file_path), 'r') as f:
            content = json.load(f)
        auth_mode = content.get('auth_mode')
        if not auth_mode or not auth_mode.strip():
            logger.error(f"No auth mode found in {file_path}")
            logger.info("Setting auth mode to default IAM")
            return AuthMode.IAM.value
        
        # Validate auth mode
        try:
            auth_mode_enum = AuthMode(auth_mode)
            logger.info(f"Returning auth mode: {auth_mode_enum.value}")
            return auth_mode_enum.value
        except ValueError:
            logger.error(f"Invalid auth mode: {auth_mode}, defaulting to IAM")
            return AuthMode.IAM.value
            
    except FileNotFoundError:
        logger.error(f"Profile not found: {file_path}")
        logger.info("Setting auth mode to default IAM")
        return AuthMode.IAM.value
    except json.JSONDecodeError:
        logger.error(f"Invalid JSON in token file: {file_path}")
        logger.info("Setting auth mode to default IAM")
        return AuthMode.IAM.value
    except Exception as e:
        logger.error(f"Error reading token file: {e}")
        logger.info("Setting auth mode to default IAM")
        return AuthMode.IAM.value

def detect_environment():
    """Detect the runtime environment using EnvironmentDetector (copied from telemetry)"""
    try:
        environment = asyncio.run(EnvironmentDetector.get_environment())
        logging.info(f"Detected environment: {environment.name}")


        return environment.name
        
    except Exception as e:
        logging.error(f"Failed to detect environment: {e}", exc_info=True)
        return "UNKNOWN"

def get_sm_environment_metadata(request_host: str) -> EnvironmentMetadata:
    """Fetch SageMaker environment metadata for origin validation.
    
    Makes HTTP request to JupyterLab API endpoint to retrieve runtime environment
    configuration including domain_id, region, and stage information.
    
    Args:
        request_host: The host from WebSocket request (e.g., 'd-abc123.notebook.us-east-1.sagemaker.aws')
        
    Returns:
        EnvironmentMetadata: Dictionary with environment configuration
        
    Raises:
        requests.RequestException: If API call fails
        ValueError: If response data is invalid or missing required fields
    """
    base_url = f"http://{request_host}"
    api_endpoint = f"{base_url}/jupyterlab/default/api/env"
    
    logger.info(f"Fetching environment metadata from: {api_endpoint}")
    
    try:
        response = requests.get(api_endpoint, timeout=5)
        response.raise_for_status()
        env_data = response.json()
        
        # Validate required fields
        required_fields = ['domain_id', 'aws_region', 'dz_stage']
        for field in required_fields:
            if field not in env_data:
                raise ValueError(f"Missing required field in environment metadata: {field}")
        
        logger.info(f"Successfully fetched environment metadata: domain_id={env_data.get('domain_id')}, region={env_data.get('aws_region')}, stage={env_data.get('dz_stage')}")
        return env_data
        
    except requests.Timeout:
        logger.error(f"Timeout fetching environment metadata from {api_endpoint}")
        raise
    except requests.RequestException as e:
        logger.error(f"Failed to fetch environment metadata: {e}")
        raise
    except (ValueError, KeyError) as e:
        logger.error(f"Invalid environment metadata response: {e}")
        raise

def build_origin_validation_patterns(env: EnvironmentMetadata) -> List[str]:
    """Build list of valid SageMaker origin patterns from environment metadata.
    
    Constructs origin URLs that match the SageMaker domain pattern:
    https://{domain_id}.sagemaker[-{stage}].{region}.on.aws
    
    Handles both domain_id and sm_domain_id (if different) to support various
    SageMaker deployment configurations.
    
    Args:
        env: Environment metadata dictionary
        
    Returns:
        List of valid origin URL strings
    """
    stage = '' if env.get('dz_stage') == 'prod' else f"-{env['dz_stage']}"
    patterns = []
    
    domain_id = env.get('domain_id')
    sm_domain_id = env.get('sm_domain_id')
    aws_region = env.get('aws_region')
    
    if domain_id and aws_region:
        pattern = f"https://{domain_id}.sagemaker{stage}.{aws_region}.on.aws"
        patterns.append(pattern)
        logger.debug(f"Added origin pattern: {pattern}")
    
    if sm_domain_id and sm_domain_id != domain_id and aws_region:
        pattern = f"https://{sm_domain_id}.sagemaker{stage}.{aws_region}.on.aws"
        patterns.append(pattern)
        logger.debug(f"Added origin pattern: {pattern}")
    
    logger.info(f"Built {len(patterns)} origin validation patterns")
    return patterns
