#!/usr/bin/env python3
import os
import sys
import json
import re
try:
    from jinja2 import Environment, FileSystemLoader
except ImportError:
    print("FATAL: jinja2 is not installed. Please run 'pip install jinja2'.", file=sys.stderr)
    exit(1)
import logging
import ipaddress
from typing import Optional

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def is_ip_address(address: str) -> bool:
    """Checks if a given string is a valid IP address (IPv4 or IPv6)."""
    if not address:
        return False
    try:
        ipaddress.ip_address(address)
        return True
    except ValueError:
        return False


def load_dotenv(dotenv_path):
    """
    Loads environment variables from a .env file.
    Does not override existing environment variables from the system.
    """
    if not os.path.exists(dotenv_path):
        logging.warning(f".env file not found at {dotenv_path}. Using system environment variables or defaults.")
        return
    try:
        with open(dotenv_path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    key = key.strip()
                    value = value.strip()
                    # Remove surrounding quotes which are common in .env files
                    if (value.startswith('"') and value.endswith('"')) or \
                       (value.startswith("'") and value.endswith("'")):
                        value = value[1:-1]
                    
                    # os.environ only takes strings
                    value = str(value)

                    if key not in os.environ:
                        os.environ[key] = value
        logging.info(f"Successfully loaded variables from {dotenv_path}")
    except Exception as e:
        logging.error(f"Failed to read or parse {dotenv_path}: {e}")
        # Continue, will use defaults or system env vars

def _get_port_from_proxy_url(url: str) -> Optional[str]:
    """Extracts the port from a proxy URL string."""
    if not url or not isinstance(url, str):
        return None
    match = re.search(r':(\d+)$', url.strip())
    return match.group(1) if match else None

def expand_env_vars(value: str) -> str:
    """
    Expands environment variables in a string, including default values.
    Supports ${VAR} and ${VAR:-default}.
    """
    if not isinstance(value, str):
        return value
    
    # Regex to find ${VAR:-default} or ${VAR}
    pattern = re.compile(r'\$\{(?P<var>\w+)(?::-(?P<default>.*?))?\}')
    
    def replacer(match):
        var_name = match.group('var')
        default_value = match.group('default')
        # Get value from os.environ, or use default, or empty string
        return os.getenv(var_name, default_value if default_value is not None else '')

    return pattern.sub(replacer, value)

def generate_configs():
    """
    Generates envoy.yaml, docker-compose.camoufox.yaml, and camoufox_endpoints.json
    from Jinja2 templates and environment variables.
    """
    try:
        # --- Load .env file ---
        script_dir = os.path.dirname(os.path.abspath(__file__))
        dotenv_path = os.path.join(script_dir, '.env')
        load_dotenv(dotenv_path)

        # --- Common Configuration ---
        ytdlp_workers_str = os.getenv('YTDLP_WORKERS', '3').strip()
        try:
            # Handle empty string case by defaulting to 3, otherwise convert to int.
            worker_count = int(ytdlp_workers_str) if ytdlp_workers_str else 3
        except (ValueError, TypeError):
            logging.warning(f"Invalid value for YTDLP_WORKERS: '{ytdlp_workers_str}'. Defaulting to 3.")
            worker_count = 3

        if worker_count == 0:
            worker_count = os.cpu_count() or 1
            logging.info(f"YTDLP_WORKERS is 0, auto-detected {worker_count} CPU cores for worker and camoufox config.")

        config_dir = os.path.join(script_dir, 'config')
        os.makedirs(config_dir, exist_ok=True)
        
        env = Environment(loader=FileSystemLoader(script_dir), trim_blocks=True, lstrip_blocks=True)

        # Get service role from environment to determine what to generate
        service_role = os.getenv('SERVICE_ROLE', 'all-in-one')
        logging.info(f"Service role for generation: '{service_role}'")

        # --- Camoufox Configuration (only for worker/all-in-one roles) ---
        camoufox_proxies = []
        expanded_camoufox_proxies_str = ""
        if service_role != 'management':
            logging.info("--- Generating Camoufox (Remote Browser) Configuration ---")
            camoufox_proxies_str = os.getenv('CAMOUFOX_PROXIES')
            if not camoufox_proxies_str:
                logging.warning("CAMOUFOX_PROXIES environment variable not set. No camoufox instances will be generated.")
            else:
                # Expand environment variables within the string before splitting
                expanded_camoufox_proxies_str = expand_env_vars(camoufox_proxies_str)
                logging.info(f"Expanded CAMOUFOX_PROXIES from '{camoufox_proxies_str}' to '{expanded_camoufox_proxies_str}'")
                camoufox_proxies = [{'url': p.strip()} for p in expanded_camoufox_proxies_str.split(',') if p.strip()]
                logging.info(f"Found {len(camoufox_proxies)} proxy/proxies for Camoufox.")
                logging.info(f"Each Camoufox instance will support {worker_count} concurrent browser sessions.")
                logging.info(f"Total browser sessions supported on this worker: {len(camoufox_proxies) * worker_count}")


            vnc_password = os.getenv('VNC_PASSWORD', 'supersecret')
            base_vnc_port = int(os.getenv('CAMOUFOX_BASE_VNC_PORT', 5901))
            camoufox_port = int(os.getenv('CAMOUFOX_PORT', 12345))
            camoufox_backend_prefix = os.getenv('CAMOUFOX_BACKEND_PREFIX', 'camoufox-')

            # --- Generate docker-compose.camoufox.yaml ---
            compose_output_file = os.path.join(script_dir, 'docker-compose.camoufox.yaml')
            
            # Generate the compose file directly without template
            with open(compose_output_file, 'w') as f:
                f.write("# THIS FILE IS AUTO-GENERATED BY generate_envoy_config.py\n")
                f.write("# DO NOT EDIT MANUALLY.\n")
                f.write("#\n")
                f.write("# It contains the service definitions for the camoufox instances\n")
                f.write("# and adds the necessary dependencies to the main services.\n")
                f.write("services:\n\n")
                
                # Generate services for each proxy
                for i, proxy in enumerate(camoufox_proxies):
                    service_name = f"camoufox-{i+1}"
                    # Each container gets its own unique range of ports to avoid conflicts
                    container_base_port = camoufox_port + i * worker_count
                    host_base_port = container_base_port

                    f.write(f"  {service_name}:\n")
                    f.write(f"    build:\n")
                    f.write(f"      context: ./camoufox\n")
                    f.write(f"      dockerfile: Dockerfile\n")
                    f.write(f"      args:\n")
                    f.write(f"        VNC_PASSWORD: {vnc_password}\n")
                    f.write(f"    image: camoufox:latest\n")
                    f.write(f"    container_name: ytdlp-ops-{service_name}-1\n")
                    f.write(f"    restart: unless-stopped\n")
                    f.write(f"    shm_size: '2gb' # Mitigates browser crashes due to shared memory limitations\n")
                    f.write(f"    ports:\n")
                    f.write(f"      - \"{host_base_port}-{host_base_port + worker_count - 1}:{container_base_port}-{container_base_port + worker_count - 1}\"\n")
                    f.write(f"    environment:\n")
                    f.write(f"      - DISPLAY=:99\n")
                    f.write(f"      - MOZ_HEADLESS_STACKSIZE=2097152\n")
                    f.write(f"      - CAMOUFOX_MAX_MEMORY_MB=2048\n")
                    f.write(f"      - CAMOUFOX_MAX_CONCURRENT_CONTEXTS=8\n")
                    f.write(f"      - CAMOUFOX_RESTART_THRESHOLD_MB=1500\n")
                    f.write(f"    volumes:\n")
                    f.write(f"      - /tmp/.X11-unix:/tmp/.X11-unix:rw\n")
                    f.write(f"      - camoufox-data-{i+1}:/app/context-data\n")
                    f.write(f"      - camoufox-browser-cache:/root/.cache/ms-playwright # Persist browser binaries\n")
                    f.write(f"    command: [\n")
                    f.write(f"      \"--ws-host\", \"0.0.0.0\",\n")
                    f.write(f"      \"--port\", \"{container_base_port}\",\n")
                    f.write(f"      \"--num-instances\", \"{worker_count}\",\n")
                    f.write(f"      \"--ws-path\", \"mypath\",\n")
                    f.write(f"      \"--proxy-url\", \"{proxy['url']}\",\n")
                    f.write(f"      \"--headless\",\n")
                    f.write(f"      \"--monitor-resources\",\n")
                    f.write(f"      \"--memory-restart-threshold\", \"1800\",\n")
                    f.write(f"      \"--preferences\", \"layers.acceleration.disabled=true,dom.ipc.processCount=2,media.memory_cache_max_size=102400,browser.cache.memory.capacity=102400\"\n")
                    f.write(f"    ]\n")
                    f.write(f"    deploy:\n")
                    f.write(f"      resources:\n")
                    f.write(f"        limits:\n")
                    f.write(f"          memory: 2.5G\n")
                    f.write(f"    logging:\n")
                    f.write(f"      driver: \"json-file\"\n")
                    f.write(f"      options:\n")
                    f.write(f"        max-size: \"100m\"\n")
                    f.write(f"        max-file: \"3\"\n")
                    f.write(f"    networks:\n")
                    f.write(f"      - proxynet\n\n")
                
                # Add camoufox-group service that depends on all camoufox instances
                if camoufox_proxies:
                    f.write("  camoufox-group:\n")
                    f.write("    image: alpine:latest\n")
                    f.write("    command: [\"echo\", \"Camoufox group ready.\"]\n")
                    f.write("    restart: \"no\"\n")
                    f.write("    depends_on:\n")
                    for i in range(len(camoufox_proxies)):
                        f.write(f"      - camoufox-{i+1}\n")
                    f.write("    networks:\n")
                    f.write("      - proxynet\n\n")
                
                # Write volumes section
                f.write("volumes:\n")
                for i in range(len(camoufox_proxies)):
                    f.write(f"  camoufox-data-{i+1}:\n")
                if camoufox_proxies:
                    f.write("  camoufox-browser-cache:\n")
                f.write("\n")
                
                # Write networks section
                f.write("networks:\n")
                f.write("  proxynet:\n")
                f.write("    name: airflow_proxynet\n")
                f.write("    external: true\n")
            
            logging.info(f"Successfully generated {compose_output_file} with {len(camoufox_proxies)} camoufox service(s).")
            logging.info("This docker-compose file defines the remote browser services, one for each proxy.")
            logging.info("----------------------------------------------------------")

            # --- Generate camoufox_endpoints.json ---
            endpoints_map = {}
            for i, proxy in enumerate(camoufox_proxies):
                proxy_port = _get_port_from_proxy_url(proxy['url'])
                if proxy_port:
                    container_base_port = camoufox_port + i * worker_count
                    endpoints = []
                    for j in range(worker_count):
                        port = container_base_port + j
                        endpoints.append(f"ws://{camoufox_backend_prefix}{i+1}:{port}/mypath")
                    
                    endpoints_map[proxy_port] = {
                        "ws_endpoints": endpoints
                    }
                else:
                    logging.warning(f"Could not extract port from proxy URL: {proxy['url']}. Skipping for endpoint map.")

            endpoints_data = {"endpoints": endpoints_map}
            camoufox_dir = os.path.join(script_dir, 'camoufox')
            endpoints_output_file = os.path.join(camoufox_dir, 'camoufox_endpoints.json')
            with open(endpoints_output_file, 'w') as f:
                json.dump(endpoints_data, f, indent=2)
            logging.info(f"Successfully generated {endpoints_output_file} with {len(endpoints_map)} port-keyed endpoint(s).")
            logging.info("This file maps each proxy to a list of WebSocket endpoints for Camoufox.")
            logging.info("The token_generator uses this map to connect to the correct remote browser.")
        else:
            logging.info("Skipping Camoufox configuration generation for 'management' role.")

        # --- Generate docker-compose-ytdlp-ops.yaml ---
        ytdlp_ops_template = env.get_template('docker-compose-ytdlp-ops.yaml.j2')
        ytdlp_ops_output_file = os.path.join(script_dir, 'docker-compose-ytdlp-ops.yaml')

        # Combine all proxies (camoufox and general) into a single string for the server.
        all_proxies = []
        if expanded_camoufox_proxies_str:
            all_proxies.extend([p.strip() for p in expanded_camoufox_proxies_str.split(',') if p.strip()])

        general_proxies_str = os.getenv('GENERAL_PROXIES')
        if general_proxies_str:
            expanded_general_proxies_str = expand_env_vars(general_proxies_str)
            logging.info(f"Expanded GENERAL_PROXIES from '{general_proxies_str}' to '{expanded_general_proxies_str}'")
            general_proxies = [p.strip() for p in expanded_general_proxies_str.split(',') if p.strip()]
            all_proxies.extend(general_proxies)
            logging.info(f"Adding {len(general_proxies)} general purpose proxy/proxies.")

        # Also check for the SOCKS5_SOCK_SERVER_IP for backward compatibility with docs
        socks_server_ip = os.getenv('SOCKS5_SOCK_SERVER_IP', '172.17.0.1')
        if socks_server_ip:
            socks_server_port = os.getenv('SOCKS5_SOCK_SERVER_PORT', '1087')
            general_proxy_url = f"socks5://{socks_server_ip}:{socks_server_port}"
            if general_proxy_url not in all_proxies:
                all_proxies.append(general_proxy_url)
                logging.info(f"Adding general purpose proxy from SOCKS5_SOCK_SERVER_IP: {general_proxy_url}")

        combined_proxies_str = ",".join(all_proxies)
        logging.info(f"Combined proxy string for ytdlp-ops-service: '{combined_proxies_str}'")

        ytdlp_ops_config_data = {
            'combined_proxies_str': combined_proxies_str,
            'service_role': service_role,
        }
        rendered_ytdlp_ops_config = ytdlp_ops_template.render(ytdlp_ops_config_data)
        with open(ytdlp_ops_output_file, 'w') as f:
            f.write(rendered_ytdlp_ops_config)
        logging.info(f"Successfully generated {ytdlp_ops_output_file}")

        # --- Envoy Configuration ---
        envoy_port = int(os.getenv('ENVOY_PORT', 9080))
        base_port = int(os.getenv('YTDLP_BASE_PORT', 9090))
        envoy_admin_port = int(os.getenv('ENVOY_ADMIN_PORT', 9901))
        # For local dev, ENVOY_BACKEND_ADDRESS is set to 127.0.0.1. For Docker, it's unset, so we default to the service name.
        backend_address = os.getenv('ENVOY_BACKEND_ADDRESS', 'ytdlp-ops-service')
        # Use STATIC for IP addresses, and STRICT_DNS for anything else (hostnames).
        envoy_cluster_type = 'STATIC' if is_ip_address(backend_address) else 'STRICT_DNS'
        
        # --- Generate envoy.yaml ---
        envoy_template = env.get_template('envoy.yaml.j2')
        envoy_output_file = os.path.join(script_dir, 'envoy.yaml')

        logging.info("--- Generating Envoy Configuration ---")
        logging.info(f"Envoy will listen on public port: {envoy_port}")
        logging.info(f"It will load balance requests across {worker_count} internal gRPC endpoints of the 'ytdlp-ops-service'.")
        logging.info(f"The backend service is located at: '{backend_address}' (type: {envoy_cluster_type})")

        envoy_config_data = {
            'envoy_port': envoy_port,
            'worker_count': worker_count,
            'base_port': base_port,
            'envoy_admin_port': envoy_admin_port,
            'backend_address': backend_address,
            'envoy_cluster_type': envoy_cluster_type,
        }
        
        rendered_envoy_config = envoy_template.render(envoy_config_data)
        with open(envoy_output_file, 'w') as f:
            f.write(rendered_envoy_config)
        logging.info(f"Successfully generated {envoy_output_file}")
        logging.info("--- Configuration Generation Complete ---")

    except Exception as e:
        logging.error(f"Failed to generate configurations: {e}", exc_info=True)
        exit(1)

if __name__ == '__main__':
    generate_configs()