"""
Authentication Module for Mesh SDK

This module handles authentication with the Mesh backend API.
"""

import os
import sys
import time
import json
import logging
import webbrowser
import urllib.parse
import requests
import random
import socket
import http.server
import socketserver
import threading
from typing import Dict, Any, Optional, Tuple, List, Union

# Import token manager functions
from .token_manager import store_token, get_token, is_token_valid, clear_token

# Import configuration
from .config import (
    get_config, 
    get_auth_config_endpoint, 
    get_auth_url_endpoint, 
    get_token_exchange_endpoint,
    get_token_refresh_endpoint
)

# Configure logging
logger = logging.getLogger("mesh.auth")

# Auth0 configuration cache
AUTH0_CONFIG_CACHE = {}

def get_auth0_config(is_device_flow=False) -> Dict[str, str]:
    """
    Get Auth0 configuration from the backend.
    
    Args:
        is_device_flow: If True, request the device flow client ID
    
    Returns:
        Dict[str, str]: Auth0 configuration including domain, client_id, and audience
    """
    global AUTH0_CONFIG_CACHE
    
    # Use a cache key that includes the flow type
    cache_key = "device" if is_device_flow else "default"
    
    # Return cached config if available and not empty
    if cache_key in AUTH0_CONFIG_CACHE and all(AUTH0_CONFIG_CACHE[cache_key].values()):
        return AUTH0_CONFIG_CACHE[cache_key]
    
    try:
        # Fetch Auth0 configuration from the backend, specifying device flow if needed
        url = get_auth_config_endpoint()
        if is_device_flow:
            url = f"{url}?flow=device"
            
        logger.debug(f"Fetching Auth0 configuration from {url}")
        response = requests.get(url)
        response.raise_for_status()
        
        config = response.json()
        if config and "domain" in config and "client_id" in config:
            # Update cache
            if not isinstance(AUTH0_CONFIG_CACHE, dict):
                AUTH0_CONFIG_CACHE = {}
            AUTH0_CONFIG_CACHE[cache_key] = config
            return config
        else:
            logger.error("Invalid Auth0 configuration received from backend")
            return {}
    except Exception as e:
        logger.error(f"Failed to fetch Auth0 configuration from backend: {str(e)}")
        return {}

def _find_available_port(start_port=8000, end_port=9000):
    """Find an available port within a range."""
    for port in range(start_port, end_port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            if s.connect_ex(('localhost', port)) != 0:
                return port
    return None

def authenticate(headless=False) -> Dict[str, Any]:
    """
    Authenticate with the Mesh backend.
    
    This function handles the authentication flow:
    1. First checks for an existing valid token in secure storage
    2. If token exists but is expired, tries to refresh it
    3. Otherwise, initiates browser-based auth or headless device flow
    
    Args:
        headless: If True, use device code flow instead of browser-based auth
        
    Returns:
        dict: Token data if authentication successful, None otherwise
    """
    # Check if we already have a valid token
    token_data = get_token()
    
    # If token exists but is invalid, try to refresh it first
    if token_data and not is_token_valid(token_data):
        logger.info("Token exists but is invalid or expired, attempting to refresh")
        try:
            # Try refreshing the token using the backend endpoint
            refreshed_token = refresh_auth_token(refresh_token=token_data.get("refresh_token"))
            
            if refreshed_token and is_token_valid(refreshed_token):
                logger.info("Successfully refreshed token")
                return refreshed_token
            else:
                logger.info("Token refresh failed, will proceed with browser-based authentication")
        except Exception as e:
            logger.warning(f"Error during token refresh: {str(e)}.")
    elif token_data and is_token_valid(token_data):
        logger.info("Using existing valid token")
        return token_data
    
    # At this point we need to authenticate
    if headless:
        return authenticate_device_flow()
    else:
        return authenticate_with_browser()

def authenticate_with_browser() -> Dict[str, Any]:
    """
    Authenticate using a browser-based flow with a temporary local HTTP server.
    
    This method:
    1. Finds an available port for the callback server
    2. Gets an auth URL from the backend
    3. Opens the browser for the user to authenticate
    4. Captures the authorization code from the callback
    5. Exchanges the code for tokens via the backend
    
    Returns:
        dict: Token data if authentication successful, None otherwise
    """
    # Find an available port for the callback server
    callback_port = _find_available_port()
    if not callback_port:
        logger.error("Could not find an available port for the callback server")
        return None
    
    logger.info(f"Using port {callback_port} for authentication callback")
    callback_uri = f"http://localhost:{callback_port}/callback"
    auth_code = {}
    
    # Define handler for the callback
    class CallbackHandler(http.server.BaseHTTPRequestHandler):
        def do_GET(self):
            parsed = urllib.parse.urlparse(self.path)
            qs = urllib.parse.parse_qs(parsed.query)
            if "code" in qs:
                auth_code["code"] = qs["code"][0]
                self.send_response(200)
                self.send_header("Content-type", "text/html")
                self.end_headers()
                self.wfile.write(b"<html><body><h1>Authentication successful. You may close this window.</h1></body></html>")
            else:
                self.send_response(400)
                self.end_headers()
        def log_message(self, format, *args):
            return  # Suppress log messages
    
    # Try to set up the callback server
    try:
        # Create a server with a random port that's available
        with socketserver.TCPServer(("", callback_port), CallbackHandler) as httpd:
            server_thread = threading.Thread(target=httpd.serve_forever)
            server_thread.daemon = True
            server_thread.start()
            
            # Generate a state value for CSRF protection
            state = f"mesh_auth_{int(time.time())}_{random.randint(1000, 9999)}"
            
            # Request an auth URL from the backend
            auth_url_data = {
                "redirect_uri": callback_uri,
                "scope": "openid profile email offline_access",
                "state": state
            }
            
            logger.info(f"Requesting auth URL using callback URI: {callback_uri}")
            try:
                auth_url_response = requests.post(get_auth_url_endpoint(), json=auth_url_data)
                auth_url_response.raise_for_status()
                auth_url_result = auth_url_response.json()
                
                if "auth_url" not in auth_url_result:
                    logger.error("No auth_url received from backend")
                    httpd.shutdown()
                    return None
                
                auth_url = auth_url_result["auth_url"]
                logger.info(f"Opening browser to: {auth_url}")
                
                # Open the browser for the user to authenticate
                try:
                    webbrowser.open(auth_url)
                    print(f"If your browser doesn't open automatically, please visit: {auth_url}")
                except Exception as e:
                    logger.warning(f"Could not open browser automatically: {str(e)}")
                    print(f"Please open this URL in your browser to authenticate: {auth_url}")
                
                # Wait for the callback with the authorization code
                timeout = 300  # 5 minutes
                poll_interval = 1
                start = time.time()
                while time.time() - start < timeout:
                    if "code" in auth_code:
                        break
                    time.sleep(poll_interval)
                
                httpd.shutdown()
                
                # If we didn't get a code, the authentication failed
                if "code" not in auth_code:
                    logger.error("Authentication timed out, did not receive code")
                    return None
                
                # Exchange the code for tokens
                code = auth_code["code"]
                logger.info("Authorization code received, exchanging for tokens")
                token_data = exchange_code_for_token(code, callback_uri)
                
                if not token_data or "access_token" not in token_data:
                    logger.error("Token exchange failed")
                    return None
                
                # Store the token and return it
                store_token(token_data)
                logger.info("Authentication successful")
                return token_data
                
            except requests.RequestException as e:
                logger.error(f"Error requesting auth URL: {str(e)}")
                httpd.shutdown()
                return None
                
    except Exception as e:
        logger.error(f"Error with authentication callback server: {str(e)}")
        return None

def exchange_code_for_token(code: str, callback_uri: str) -> Dict[str, Any]:
    """
    Exchange authorization code for token using the backend endpoint.

    Args:
        code: Authorization code from Auth0
        callback_uri: The callback URI that was used to get the code
        
    Returns:
        dict: Token data or empty dict if exchange failed
    """
    try:
        exchange_url = get_token_exchange_endpoint()
        exchange_data = {
            "code": code,
            "redirect_uri": callback_uri
        }
        
        logger.debug(f"Exchanging code for token at: {exchange_url}")
        response = requests.post(exchange_url, json=exchange_data)
        response.raise_for_status()
        
        token_data = response.json()
        
        # Add expires_at field if not present
        if "expires_in" in token_data and "expires_at" not in token_data:
            token_data["expires_at"] = int(time.time()) + token_data["expires_in"]
            
        # Store the token
        store_token(token_data)
        return token_data
    except Exception as e:
        logger.error(f"Error exchanging code for token: {str(e)}")
        return {}

def refresh_auth_token(refresh_token=None):
    """
    Refresh an Auth0 token using the refresh token.
    
    Args:
        refresh_token: Refresh token to use
        
    Returns:
        dict: New token data or None if refresh failed
    """
    # If no refresh token was provided, try to get it from the stored token
    if not refresh_token:
        token_data = get_token()
        if token_data:
            refresh_token = token_data.get("refresh_token")
    
    # If we still don't have a refresh token, we can't refresh
    if not refresh_token:
        logger.warning("No refresh token available")
        return None
    
    # If the refresh token is empty, we can't refresh
    if not refresh_token.strip():
        logger.warning("Refresh token is empty")
        return None
    
    logger.info("Attempting to refresh token")
    
    try:
        # Try to refresh using the backend
        response = requests.post(
            get_token_refresh_endpoint(),
            json={"refresh_token": refresh_token},
            headers={"Content-Type": "application/json"}
        )
        
        if response.status_code == 200:
            token_data = response.json()
            
            # Add expires_at for convenience
            if "expires_in" in token_data:
                token_data["expires_at"] = int(time.time()) + token_data["expires_in"]
            
            # Store the new token
            store_token(token_data)
            
            logger.info("Successfully refreshed token")
            return token_data
        else:
            logger.warning(f"Token refresh failed: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        logger.warning(f"Token refresh failed: {str(e)}")
        return None

def authenticate_device_flow():
    """
    Authenticate using device code flow.
    
    This is useful for headless environments, CLI tools, or notebook environments like Colab.
    
    Returns:
        dict: Token data if authentication successful, None otherwise
    """
    try:
        # Check if running in Jupyter/IPython environment
        in_notebook = False
        try:
            from IPython import get_ipython
            if get_ipython() is not None:
                in_notebook = True
        except ImportError:
            pass
        
        # First get the auth0 config from the backend, specifying this is for device flow
        auth0_config = get_auth0_config(is_device_flow=True)
        if not auth0_config:
            logger.error("Could not get Auth0 configuration")
            return None
        
        # Get a device code from Auth0
        device_code_url = f"https://{auth0_config['domain']}/oauth/device/code"
        device_code_data = {
            "client_id": auth0_config['client_id'],
            "scope": "openid profile email offline_access",
            "audience": auth0_config.get('audience', '')
        }
        
        device_response = requests.post(device_code_url, data=device_code_data)
        device_response.raise_for_status()
        device_result = device_response.json()
        
        if "verification_uri_complete" not in device_result or "device_code" not in device_result:
            logger.error("Invalid device code response")
            return None
        
        verification_uri = device_result["verification_uri_complete"]
        user_code = device_result.get("user_code", "")
        device_code = device_result["device_code"]
        polling_interval = device_result.get("interval", 5)
        expiry = device_result.get("expires_in", 900)  # Default to 15 minutes
        
        # Display the verification URI for the user, with enhanced rendering for notebooks
        if in_notebook:
            from IPython.display import display, HTML, Markdown
            
            html_content = f"""
            <div style="background-color: #f8f9fa; padding: 15px; border-radius: 5px; margin: 10px 0; border: 1px solid #dee2e6;">
                <h3 style="margin-top: 0;">Mesh SDK Authentication</h3>
                <p>To authenticate the Mesh SDK, please click the button below or visit this URL on any device:</p>
                <div style="margin: 15px 0;">
                    <a href="{verification_uri}" target="_blank" style="display: inline-block; background-color: #4285f4; color: white; padding: 8px 16px; text-decoration: none; border-radius: 4px; font-weight: bold;">
                        Authenticate with code: {user_code}
                    </a>
                </div>
                <p style="margin-bottom: 0; font-size: 90%; color: #6c757d;">This code will expire in {expiry} seconds.</p>
            </div>
            """
            display(HTML(html_content))
            
            # Also show waiting message
            display(Markdown("*Waiting for authentication... You can continue working in other cells while this completes.*"))
        else:
            # Standard console output for CLI environments
            print(f"\nTo authenticate the Mesh SDK, please visit this URL on any device:")
            print(f"\n{verification_uri}\n")
            if user_code:
                print(f"Your verification code is: {user_code}")
            print(f"This code will expire in {expiry} seconds.\n")
        
        # Poll for token
        token_url = f"https://{auth0_config['domain']}/oauth/token"
        start_time = time.time()
        
        if in_notebook:
            from IPython.display import clear_output
        
        while time.time() - start_time < expiry:
            try:
                token_response = requests.post(token_url, data={
                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                    "device_code": device_code,
                    "client_id": auth0_config['client_id']
                })
                
                if token_response.status_code == 200:
                    token_data = token_response.json()
                    if "access_token" not in token_data:
                        logger.error("Token response missing access_token")
                        return None
                    
                    # Add expires_at for convenience
                    if "expires_in" in token_data:
                        token_data["expires_at"] = int(time.time()) + token_data["expires_in"]
                    
                    # Store the token
                    store_token(token_data)
                    
                    if in_notebook:
                        clear_output(wait=True)
                        display(HTML("""
                        <div style="background-color: #d4edda; padding: 15px; border-radius: 5px; margin: 10px 0; border: 1px solid #c3e6cb;">
                            <h3 style="margin-top: 0; color: #155724;">Authentication Successful!</h3>
                            <p style="margin-bottom: 0; color: #155724;">You can now use the Mesh SDK.</p>
                        </div>
                        """))
                    else:
                        print("\nAuthentication successful! You can now use the Mesh SDK.")
                    
                    return token_data
                elif token_response.status_code == 403:
                    # User has not yet authorized
                    time.sleep(polling_interval)
                else:
                    error_data = token_response.json()
                    error = error_data.get("error", "unknown_error")
                    
                    if error == "expired_token":
                        logger.error("Device code expired")
                        if in_notebook:
                            clear_output(wait=True)
                            display(HTML("""
                            <div style="background-color: #f8d7da; padding: 15px; border-radius: 5px; margin: 10px 0; border: 1px solid #f5c6cb;">
                                <h3 style="margin-top: 0; color: #721c24;">Authentication Failed</h3>
                                <p style="margin-bottom: 0; color: #721c24;">The device code has expired. Please try again.</p>
                            </div>
                            """))
                        else:
                            print("\nThe device code has expired. Please try again.")
                        return None
                    elif error == "access_denied":
                        logger.error("User denied access")
                        if in_notebook:
                            clear_output(wait=True)
                            display(HTML("""
                            <div style="background-color: #f8d7da; padding: 15px; border-radius: 5px; margin: 10px 0; border: 1px solid #f5c6cb;">
                                <h3 style="margin-top: 0; color: #721c24;">Authentication Denied</h3>
                                <p style="margin-bottom: 0; color: #721c24;">Authentication was denied. Please try again.</p>
                            </div>
                            """))
                        else:
                            print("\nAuthentication was denied. Please try again.")
                        return None
                    else:
                        time.sleep(polling_interval)
            except Exception as e:
                logger.error(f"Error polling for token: {str(e)}")
                time.sleep(polling_interval)
        
        if in_notebook:
            clear_output(wait=True)
            display(HTML("""
            <div style="background-color: #f8d7da; padding: 15px; border-radius: 5px; margin: 10px 0; border: 1px solid #f5c6cb;">
                <h3 style="margin-top: 0; color: #721c24;">Authentication Timed Out</h3>
                <p style="margin-bottom: 0; color: #721c24;">Authentication timed out. Please try again.</p>
            </div>
            """))
        else:
            print("\nAuthentication timed out. Please try again.")
        return None
        
    except Exception as e:
        logger.error(f"Error during device flow authentication: {str(e)}")
        return None