"""OAuth authentication and token refresh with Dynamic Client Registration."""

import webbrowser
import socket
import os
from http.server import HTTPServer, BaseHTTPRequestHandler
from datetime import datetime, timedelta
import requests
from urllib.parse import urlparse, parse_qs
from mcpbundles_proxy.config import save_token, load_token
import secrets
import hashlib
import base64

# Load environment variables
try:
    from dotenv import load_dotenv  # type: ignore
    load_dotenv()
except ImportError:
    pass


def find_free_port(start_port=8765, max_attempts=10):
    """Find available port starting from start_port."""
    for port in range(start_port, start_port + max_attempts):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('localhost', port))
                return port
        except OSError:
            continue
    raise RuntimeError("Could not find available port for OAuth callback")


class OAuthCallbackHandler(BaseHTTPRequestHandler):
    """Handle OAuth callback from browser."""
    
    auth_code: str | None = None
    
    def do_GET(self):
        """Handle GET request with OAuth code."""
        query = parse_qs(urlparse(self.path).query)
        OAuthCallbackHandler.auth_code = query.get('code', [None])[0]
        
        # Show success page in browser
        self.send_response(200)
        self.send_header('Content-type', 'text/html')
        self.end_headers()
        html = """
            <html>
            <head>
                <style>
                    body {
                        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
                        text-align: center;
                        padding: 50px;
                        background: #f5f5f5;
                    }
                    .container {
                        background: white;
                        border-radius: 10px;
                        padding: 40px;
                        max-width: 500px;
                        margin: 0 auto;
                        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
                    }
                    h1 { color: #28a745; }
                </style>
            </head>
            <body>
                <div class="container">
                    <h1>&#x2705; Authentication Successful</h1>
                    <p>You can close this window and return to your terminal.</p>
                </div>
            </body>
            </html>
        """
        self.wfile.write(html.encode('utf-8'))
    
    def log_message(self, format, *args):
        """Suppress HTTP server logs."""
        pass


def register_client():
    """Register new OAuth client via Dynamic Client Registration (DCR)."""
    print("Registering new client with MCPBundles...")
    
    # Register multiple callback ports (8765-8774) to handle port conflicts
    redirect_uris = [f"http://localhost:{port}/callback" for port in range(8765, 8775)]
    
    # Use production URL with dev override
    register_url = "https://api.mcpbundles.com/o/register/"
    if os.getenv('ENVIRONMENT') == 'development':
        register_url = "http://localhost:8000/o/register/"
    
    response = requests.post(
        register_url,
        json={
            "client_name": "MCPBundles Desktop Proxy",
            "redirect_uris": redirect_uris,
            "grant_types": ["authorization_code", "refresh_token"],
            "response_types": ["code"],
            "scope": "read write",
            "token_endpoint_auth_method": "client_secret_post"
        }
    )
    
    if response.status_code != 201:
        raise Exception(f"Client registration failed: {response.text}")
    
    client_data = response.json()
    print(f"✅ Client registered: {client_data['client_id']}")
    
    return client_data['client_id'], client_data['client_secret']


def get_or_register_client():
    """Get existing client credentials or register new client."""
    # Check if we already have credentials stored
    token_data = load_token()
    if token_data and 'client_id' in token_data and 'client_secret' in token_data:
        return token_data['client_id'], token_data['client_secret']
    
    # Register new client
    return register_client()


def generate_pkce():
    """Generate PKCE code_verifier and code_challenge."""
    # Generate random code_verifier (43-128 characters)
    code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode('utf-8').rstrip('=')
    
    # Generate code_challenge (SHA256 hash of code_verifier)
    challenge_bytes = hashlib.sha256(code_verifier.encode('utf-8')).digest()
    code_challenge = base64.urlsafe_b64encode(challenge_bytes).decode('utf-8').rstrip('=')
    
    return code_verifier, code_challenge


def oauth_flow():
    """Execute OAuth flow with browser and local callback server."""
    # Get or register OAuth client
    client_id, client_secret = get_or_register_client()
    
    # Generate PKCE parameters
    code_verifier, code_challenge = generate_pkce()
    
    # Find available port
    port = find_free_port()
    redirect_uri = f"http://localhost:{port}/callback"
    
    # Start local callback server
    server = HTTPServer(('localhost', port), OAuthCallbackHandler)
    
    # Open browser to OAuth page with PKCE
    auth_url = (
        f"https://api.mcpbundles.com/o/authorize/"
        f"?client_id={client_id}"
        f"&redirect_uri={redirect_uri}"
        f"&response_type=code"
        f"&scope=read write"
        f"&code_challenge={code_challenge}"
        f"&code_challenge_method=S256"
    )
    
    # Allow development override
    if os.getenv('ENVIRONMENT') == 'development':
        auth_url = auth_url.replace("https://api.mcpbundles.com", "http://localhost:8000")
    
    print(f"Opening browser for authentication...")
    print(f"If browser doesn't open, visit: {auth_url}")
    webbrowser.open(auth_url)
    
    # Wait for callback (with timeout)
    server.timeout = 120  # 2 minute timeout
    server.handle_request()
    
    if not OAuthCallbackHandler.auth_code:
        raise Exception("Authentication failed or timed out")
    
    # Exchange code for tokens
    print("Exchanging authorization code for tokens...")
    token_url = "https://api.mcpbundles.com/o/token/"
    if os.getenv('ENVIRONMENT') == 'development':
        token_url = "http://localhost:8000/o/token/"
    
    response = requests.post(
        token_url,
        data={
            "grant_type": "authorization_code",
            "code": OAuthCallbackHandler.auth_code,
            "client_id": client_id,
            "client_secret": client_secret,
            "redirect_uri": redirect_uri,
            "code_verifier": code_verifier
        }
    )
    
    if response.status_code != 200:
        raise Exception(f"Token exchange failed: {response.text}")
    
    token_data = response.json()
    
    # Add expiry timestamp
    expires_in = token_data.get('expires_in', 3600)
    token_data['expires_at'] = (
        datetime.now() + timedelta(seconds=expires_in)
    ).isoformat()
    
    # Store client credentials with tokens
    token_data['client_id'] = client_id
    token_data['client_secret'] = client_secret
    
    # Save tokens
    save_token(token_data)
    print(f"✅ Successfully authenticated!")
    
    return token_data


async def refresh_token_if_needed(token_data):
    """Check and refresh access token if needed."""
    expires_at = datetime.fromisoformat(token_data['expires_at'])
    
    # Refresh if expires in less than 1 hour
    if datetime.now() >= expires_at - timedelta(hours=1):
        try:
            print("Refreshing access token...")
            
            # Use production URL (allow dev override)
            token_url = "https://api.mcpbundles.com/o/token/"
            if os.getenv('ENVIRONMENT') == 'development':
                token_url = "http://localhost:8000/o/token/"
            
            response = requests.post(
                token_url,
                data={
                    "grant_type": "refresh_token",
                    "refresh_token": token_data['refresh_token'],
                    "client_id": token_data['client_id'],
                    "client_secret": token_data['client_secret']
                }
            )
            
            if response.status_code == 200:
                new_token = response.json()
                expires_in = new_token.get('expires_in', 3600)
                new_token['expires_at'] = (
                    datetime.now() + timedelta(seconds=expires_in)
                ).isoformat()
                # Keep client credentials and refresh token
                new_token['client_id'] = token_data['client_id']
                new_token['client_secret'] = token_data['client_secret']
                if 'refresh_token' not in new_token:
                    new_token['refresh_token'] = token_data['refresh_token']
                save_token(new_token)
                print("✅ Token refreshed")
                return new_token
            else:
                # Refresh failed, need to re-authenticate
                print("⚠️ Token refresh failed, re-authentication needed")
                return None
        except Exception as e:
            print(f"❌ Token refresh error: {e}")
            return None
    
    return token_data

