#!/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:
        # --- Setup Paths ---
        # The script runs from /app. Configs and templates are in /app/configs.
        project_root = os.path.dirname(os.path.abspath(__file__)) # This will be /app
        configs_dir = os.path.join(project_root, 'configs')
        
        # Load .env from the 'configs' directory.
        dotenv_path = os.path.join(configs_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.")

        # The templates are in the 'configs' directory.
        env = Environment(loader=FileSystemLoader(configs_dir), trim_blocks=True, lstrip_blocks=True)
        # Make the helper function available to Jinja2 templates
        env.globals['_get_port_from_proxy_url'] = _get_port_from_proxy_url

        # 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_template = env.get_template('docker-compose.camoufox.yaml.j2')
            compose_output_file = os.path.join(configs_dir, 'docker-compose.camoufox.yaml')

            camoufox_config_data = {
                'camoufox_proxies': camoufox_proxies,
                'vnc_password': vnc_password,
                'camoufox_port': camoufox_port,
                'worker_count': worker_count,
            }

            rendered_compose_config = compose_template.render(camoufox_config_data)
            with open(compose_output_file, 'w') as f:
                f.write(rendered_compose_config)

            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:
                    # Use the correct container name pattern that matches the docker-compose template
                    # The container name in the template is: ytdlp-ops-camoufox-{{ proxy_port }}-{{ loop.index }}-1
                    container_name = f"ytdlp-ops-camoufox-{proxy_port}-{i+1}-1"
                    container_base_port = camoufox_port + i * worker_count
                    endpoints = []
                    for j in range(worker_count):
                        port = container_base_port + j
                        endpoints.append(f"ws://{container_name}:{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}
            # The camoufox directory is at the root of the project context, not under 'airflow'.
            # camoufox_dir = os.path.join(project_root, 'camoufox')
            # os.makedirs(camoufox_dir, exist_ok=True)
            endpoints_output_file = os.path.join(configs_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(configs_dir, 'docker-compose-ytdlp-ops.yaml')

        # Combine all proxies (camoufox and general) into a single string for the server.
        all_proxies = []
        
        # Track if we have any explicit proxy configuration
        has_explicit_proxies = False
        
        # Add camoufox proxies if they exist
        if expanded_camoufox_proxies_str:
            camoufox_proxy_list = [p.strip() for p in expanded_camoufox_proxies_str.split(',') if p.strip()]
            all_proxies.extend(camoufox_proxy_list)
            if camoufox_proxy_list:
                has_explicit_proxies = True
                logging.info(f"Added {len(camoufox_proxy_list)} camoufox proxies: {camoufox_proxy_list}")


        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,
            'camoufox_proxies': camoufox_proxies,
        }
        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')
        # Output envoy.yaml to the configs directory, where other generated files are.
        envoy_output_file = os.path.join(configs_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()