import boto3
from botocore.exceptions import ClientError
import base64
import logging
import json
from typing import Optional

from desktop_env.providers.base import Provider
from desktop_env.providers.aws.proxy_pool import get_global_proxy_pool, init_proxy_pool, ProxyInfo

logger = logging.getLogger("desktopenv.providers.aws.AWSProviderWithProxy")
logger.setLevel(logging.INFO)

WAIT_DELAY = 15
MAX_ATTEMPTS = 10


class AWSProviderWithProxy(Provider):
    
    def __init__(self, region: str = None, proxy_config_file: str = None):
        super().__init__(region)
        self.current_proxy: Optional[ProxyInfo] = None
        
        # 初始化代理池
        if proxy_config_file:
            init_proxy_pool(proxy_config_file)
            logger.info(f"Initialized proxy pool from {proxy_config_file}")
        
        # 获取下一个可用代理
        self._rotate_proxy()

    def _rotate_proxy(self):
        """轮换到下一个可用代理"""
        proxy_pool = get_global_proxy_pool()
        self.current_proxy = proxy_pool.get_next_proxy()
        
        if self.current_proxy:
            logger.info(f"Switched to proxy: {self.current_proxy.host}:{self.current_proxy.port}")
        else:
            logger.warning("No proxy available, using direct connection")

    def _generate_proxy_user_data(self) -> str:
        """生成包含代理配置的user data脚本"""
        if not self.current_proxy:
            return ""
        
        proxy_url = self._format_proxy_url(self.current_proxy)
        
        user_data_script = f"""#!/bin/bash
# Configure system proxy
echo 'export http_proxy={proxy_url}' >> /etc/environment
echo 'export https_proxy={proxy_url}' >> /etc/environment
echo 'export HTTP_PROXY={proxy_url}' >> /etc/environment  
echo 'export HTTPS_PROXY={proxy_url}' >> /etc/environment

# Configure apt proxy
cat > /etc/apt/apt.conf.d/95proxy << EOF
Acquire::http::Proxy "{proxy_url}";
Acquire::https::Proxy "{proxy_url}";
EOF

# Configure chrome/chromium proxy
mkdir -p /etc/opt/chrome/policies/managed
cat > /etc/opt/chrome/policies/managed/proxy.json << EOF
{{
    "ProxyMode": "fixed_servers",
    "ProxyServer": "{self.current_proxy.host}:{self.current_proxy.port}"
}}
EOF

# Configure chromium proxy (Ubuntu default)
mkdir -p /etc/chromium/policies/managed
cat > /etc/chromium/policies/managed/proxy.json << EOF
{{
    "ProxyMode": "fixed_servers", 
    "ProxyServer": "{self.current_proxy.host}:{self.current_proxy.port}"
}}
EOF

# Configure firefox proxy - support multiple possible paths
for firefox_dir in /etc/firefox/policies /usr/lib/firefox/distribution/policies /etc/firefox-esr/policies; do
    if [ -d "$(dirname "$firefox_dir")" ]; then
        mkdir -p "$firefox_dir"
        cat > "$firefox_dir/policies.json" << EOF
{{
    "policies": {{
        "Proxy": {{
            "Mode": "manual",
            "HTTPProxy": "{self.current_proxy.host}:{self.current_proxy.port}",
            "HTTPSProxy": "{self.current_proxy.host}:{self.current_proxy.port}",
            "UseHTTPProxyForAllProtocols": true
        }}
    }}
}}
EOF
        break
    fi
done

# Reload environment variables
source /etc/environment

# Log proxy configuration
echo "$(date): Configured proxy {self.current_proxy.host}:{self.current_proxy.port}" >> /var/log/proxy-setup.log
"""
        
        return base64.b64encode(user_data_script.encode()).decode()

    def _format_proxy_url(self, proxy: ProxyInfo) -> str:
        """格式化代理URL"""
        if proxy.username and proxy.password:
            return f"{proxy.protocol}://{proxy.username}:{proxy.password}@{proxy.host}:{proxy.port}"
        else:
            return f"{proxy.protocol}://{proxy.host}:{proxy.port}"

    def start_emulator(self, path_to_vm: str, headless: bool, *args, **kwargs):
        logger.info("Starting AWS VM with proxy configuration...")
        ec2_client = boto3.client('ec2', region_name=self.region)

        try:
            # 如果实例已经存在，直接启动
            ec2_client.start_instances(InstanceIds=[path_to_vm])
            logger.info(f"Instance {path_to_vm} is starting...")

            # Wait for the instance to be in the 'running' state
            waiter = ec2_client.get_waiter('instance_running')
            waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS})
            logger.info(f"Instance {path_to_vm} is now running.")

        except ClientError as e:
            logger.error(f"Failed to start the AWS VM {path_to_vm}: {str(e)}")
            raise

    def create_instance_with_proxy(self, image_id: str, instance_type: str, 
                                 security_groups: list, subnet_id: str) -> str:
        """创建带有代理配置的新实例"""
        ec2_client = boto3.client('ec2', region_name=self.region)
        
        user_data = self._generate_proxy_user_data()
        
        run_instances_params = {
            "MaxCount": 1,
            "MinCount": 1,
            "ImageId": image_id,
            "InstanceType": instance_type,
            "EbsOptimized": True,
            "NetworkInterfaces": [
                {
                    "SubnetId": subnet_id,
                    "AssociatePublicIpAddress": True,
                    "DeviceIndex": 0,
                    "Groups": security_groups
                }
            ]
        }
        
        if user_data:
            run_instances_params["UserData"] = user_data
            
        try:
            response = ec2_client.run_instances(**run_instances_params)
            instance_id = response['Instances'][0]['InstanceId']
            
            logger.info(f"Created new instance {instance_id} with proxy configuration")
            
            logger.info(f"Waiting for instance {instance_id} to be running...")
            ec2_client.get_waiter('instance_running').wait(InstanceIds=[instance_id])
            logger.info(f"Instance {instance_id} is ready.")

            try:
                instance_details = ec2_client.describe_instances(InstanceIds=[instance_id])
                instance = instance_details['Reservations'][0]['Instances'][0]
                public_ip = instance.get('PublicIpAddress', '')
                if public_ip:
                    vnc_url = f"http://{public_ip}:5910/vnc.html"
                    logger.info("="*80)
                    logger.info(f"🖥️  VNC Web Access URL: {vnc_url}")
                    logger.info(f"📡 Public IP: {public_ip}")
                    logger.info(f"🆔 Instance ID: {instance_id}")
                    if self.current_proxy:
                        logger.info(f"🌐 Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
                    logger.info("="*80)
                    print(f"\n🌐 VNC Web Access URL: {vnc_url}")
                    if self.current_proxy:
                        print(f"🔄 Current Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
                    print(f"📍 Please open the above address in the browser for remote desktop access\n")
            except Exception as e:
                logger.warning(f"Failed to get VNC address for instance {instance_id}: {e}")
            
            return instance_id
            
        except ClientError as e:
            logger.error(f"Failed to create instance with proxy: {str(e)}")
            if self.current_proxy:
                proxy_pool = get_global_proxy_pool()
                proxy_pool.mark_proxy_failed(self.current_proxy)
                self._rotate_proxy()
            raise

    def get_ip_address(self, path_to_vm: str) -> str:
        logger.info("Getting AWS VM IP address...")
        ec2_client = boto3.client('ec2', region_name=self.region)

        try:
            response = ec2_client.describe_instances(InstanceIds=[path_to_vm])
            for reservation in response['Reservations']:
                for instance in reservation['Instances']:
                    private_ip_address = instance.get('PrivateIpAddress', '')
                    public_ip_address = instance.get('PublicIpAddress', '')

                    if public_ip_address:
                        vnc_url = f"http://{public_ip_address}:5910/vnc.html"
                        logger.info("="*80)
                        logger.info(f"🖥️  VNC Web Access URL: {vnc_url}")
                        logger.info(f"📡 Public IP: {public_ip_address}")
                        logger.info(f"🏠 Private IP: {private_ip_address}")
                        if self.current_proxy:
                            logger.info(f"🌐 Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
                        logger.info("="*80)
                        print(f"\n🌐 VNC Web Access URL: {vnc_url}")
                        if self.current_proxy:
                            print(f"🔄 Current Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
                        print(f"📍 Please open the above address in the browser for remote desktop access\n")
                    else:
                        logger.warning("No public IP address available for VNC access")
                    
                    return private_ip_address
            return ''
        except ClientError as e:
            logger.error(f"Failed to retrieve IP address for the instance {path_to_vm}: {str(e)}")
            raise

    def save_state(self, path_to_vm: str, snapshot_name: str):
        logger.info("Saving AWS VM state...")
        ec2_client = boto3.client('ec2', region_name=self.region)

        try:
            image_response = ec2_client.create_image(InstanceId=path_to_vm, Name=snapshot_name)
            image_id = image_response['ImageId']
            logger.info(f"AMI {image_id} created successfully from instance {path_to_vm}.")
            return image_id
        except ClientError as e:
            logger.error(f"Failed to create AMI from the instance {path_to_vm}: {str(e)}")
            raise

    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
        logger.info(f"Reverting AWS VM to snapshot: {snapshot_name}...")
        ec2_client = boto3.client('ec2', region_name=self.region)

        try:
            # Get original instance details for config.
            instance_details = ec2_client.describe_instances(InstanceIds=[path_to_vm])
            instance = instance_details['Reservations'][0]['Instances'][0]
            security_groups = [sg['GroupId'] for sg in instance['SecurityGroups']]
            subnet_id = instance['SubnetId']
            instance_type = instance['InstanceType']

            # Terminate the old instance. This is a non-blocking call.
            logger.info(f"Initiating termination for old instance {path_to_vm}...")
            ec2_client.terminate_instances(InstanceIds=[path_to_vm])
            logger.info(f"Old instance {path_to_vm} termination initiated.")

            # Rotate to a new proxy
            self._rotate_proxy()
            
            # Create a new instance
            new_instance_id = self.create_instance_with_proxy(
                snapshot_name, instance_type, security_groups, subnet_id
            )
            
            # Note: VNC address is displayed within create_instance_with_proxy
            logger.info(f"Successfully launched new instance {new_instance_id} for revert.")

            return new_instance_id

        except ClientError as e:
            logger.error(f"Failed to revert to snapshot {snapshot_name} for the instance {path_to_vm}: {str(e)}")
            raise

    def stop_emulator(self, path_to_vm, region=None):
        logger.info(f"Stopping AWS VM {path_to_vm}...")
        ec2_client = boto3.client('ec2', region_name=self.region)

        try:
            ec2_client.stop_instances(InstanceIds=[path_to_vm])
            waiter = ec2_client.get_waiter('instance_stopped')
            waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS})
            logger.info(f"Instance {path_to_vm} has been stopped.")
        except ClientError as e:
            logger.error(f"Failed to stop the AWS VM {path_to_vm}: {str(e)}")
            raise

    def get_current_proxy_info(self) -> Optional[dict]:
        """获取当前代理信息"""
        if self.current_proxy:
            return {
                'host': self.current_proxy.host,
                'port': self.current_proxy.port,
                'protocol': self.current_proxy.protocol,
                'failed_count': self.current_proxy.failed_count
            }
        return None

    def force_rotate_proxy(self):
        """强制轮换代理"""
        logger.info("Force rotating proxy...")
        if self.current_proxy:
            proxy_pool = get_global_proxy_pool()
            proxy_pool.mark_proxy_failed(self.current_proxy)
        self._rotate_proxy()

    def get_proxy_stats(self) -> dict:
        """获取代理池统计信息"""
        proxy_pool = get_global_proxy_pool()
        return proxy_pool.get_stats() 