"""
quiXzoom SSO Authentication Service
Microsoft-like seamless auth across all quixzoom properties

Architecture:
- Shared JWT tokens signed with RS256
- Cookie domain: .quixzoom.com (shared across all subdomains)
- Access token: short-lived (15 min)
- Refresh token: long-lived (30 days), rotated on each use
- Cross-domain: all quixzoom.* domains trust the same auth service
"""

import os
import time
import uuid
import hashlib
import secrets
from datetime import datetime, timedelta, timezone
from typing import Optional, Dict, Any
from fastapi import FastAPI, Request, Response, HTTPException, Cookie, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel, Field
import jwt
import redis
import uvicorn

app = FastAPI(
    title="quiXzoom SSO Auth",
    description="Single Sign-On for all quixzoom properties",
    version="2.0.0"
)

# Redis for sessions and refresh tokens
redis_client = redis.Redis(
    host=os.getenv('REDIS_HOST', 'localhost'),
    port=int(os.getenv('REDIS_PORT', 6379)),
    db=int(os.getenv('REDIS_DB', 0)),
    decode_responses=True
)

# JWT Configuration - load from files
# Use quixzoom- prefixed keys (actual filenames)
JWT_PRIVATE_KEY_PATH = os.getenv('JWT_PRIVATE_KEY_PATH', '/app/keys/quixzoom-private.pem')
JWT_PUBLIC_KEY_PATH = os.getenv('JWT_PUBLIC_KEY_PATH', '/app/keys/quixzoom-public.pem')

def load_key(path: str, key_type: str = 'private') -> str:
    """Load key from file, trying variations of the filename"""
    # Try the exact path first
    paths_to_try = [path]
    
    # If path ends with private.pem (without quixzoom- prefix), try with prefix
    if '/private.pem' in path and 'quixzoom' not in path:
        paths_to_try.append(path.replace('/private.pem', '/quixzoom-private.pem'))
    if '/public.pem' in path and 'quixzoom' not in path:
        paths_to_try.append(path.replace('/public.pem', '/quixzoom-public.pem'))
    
    # Also try the default names in /app/keys
    paths_to_try.append('/app/keys/quixzoom-private.pem')
    paths_to_try.append('/app/keys/quixzoom-public.pem')
    
    for p in paths_to_try:
        try:
            with open(p, 'r') as f:
                content = f.read()
                print(f"Loaded key from: {p} ({len(content)} chars)", flush=True)
                return content
        except FileNotFoundError:
            continue
    
    # Fallback to env var
    fallback = os.getenv('JWT_PRIVATE_KEY' if key_type == 'private' else 'JWT_PUBLIC_KEY', '')
    if fallback:
        print(f"Using fallback env var for {key_type} key", flush=True)
    return fallback

JWT_PRIVATE_KEY = load_key(JWT_PRIVATE_KEY_PATH, 'private')
JWT_PUBLIC_KEY = load_key(JWT_PUBLIC_KEY_PATH, 'public')
JWT_ALGORITHM = 'RS256'

print(f"Final - Private key: {len(JWT_PRIVATE_KEY)} chars, Public key: {len(JWT_PUBLIC_KEY)} chars", flush=True)
ACCESS_TOKEN_TTL = 900  # 15 minutes
REFRESH_TOKEN_TTL = 2592000  # 30 days

# Cookie settings
COOKIE_DOMAIN = '.quixzoom.com'
COOKIE_SECURE = True
COOKIE_SAMESITE = 'lax'

# Allowed origins - all quixzoom properties
ALLOWED_ORIGINS = [
    'https://quixzoom.com',
    'https://www.quixzoom.com',
    'https://app.quixzoom.com',
    'https://quixzoom.se',
    'https://www.quixzoom.se',
    'https://quixzoom.de',
    'https://www.quixzoom.de',
    'https://quixzoom.fr',
    'https://www.quixzoom.fr',
    'https://quixzoom.nl',
    'https://www.quixzoom.nl',
    'https://quixzoom.co.uk',
    'https://www.quixzoom.co.uk',
    'https://quixzoom.asia',
    'https://www.quixzoom.asia',
    'https://quixzoom.es',
    'https://www.quixzoom.es',
    'https://quixzoom.it',
    'https://www.quixzoom.it',
    'https://quixzoom.pl',
    'https://www.quixzoom.pl',
    'http://localhost:3000',
    'http://localhost:8080',
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=['*'],
)

# ─── Models ──────────────────────────────────────────────────────────────

class LoginRequest(BaseModel):
    email: str
    password: str
    device_id: Optional[str] = None
    device_info: Optional[Dict[str, Any]] = None
    redirect_url: Optional[str] = 'https://www.quixzoom.com/mina-sidor'

class RegisterRequest(BaseModel):
    email: str
    password: str
    first_name: str
    last_name: str
    phone: Optional[str] = None
    country: Optional[str] = None
    device_id: Optional[str] = None
    device_info: Optional[Dict[str, Any]] = None

class TokenResponse(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str = 'Bearer'
    expires_in: int = ACCESS_TOKEN_TTL
    user: Dict[str, Any]

class UserProfile(BaseModel):
    id: str
    email: str
    first_name: str
    last_name: str
    avatar_url: Optional[str] = None
    role: str = 'zoomer'
    kyc_status: str = 'pending'
    wallet_balance: float = 0.0
    qz_tokens: float = 0.0

# ─── Helper Functions ────────────────────────────────────────────────────

def generate_token_id() -> str:
    return secrets.token_urlsafe(32)

def create_access_token(user_id: str, email: str, role: str, token_id: str) -> str:
    now = datetime.now(timezone.utc)
    payload = {
        'sub': user_id,
        'email': email,
        'role': role,
        'jti': token_id,
        'iat': now,
        'exp': now + timedelta(seconds=ACCESS_TOKEN_TTL),
        'type': 'access',
        'iss': 'https://auth.quixzoom.com',
        'aud': 'quixzoom-platform',
    }
    return jwt.encode(payload, JWT_PRIVATE_KEY, algorithm=JWT_ALGORITHM)

def create_refresh_token(user_id: str, token_id: str, device_id: Optional[str] = None) -> str:
    now = datetime.now(timezone.utc)
    payload = {
        'sub': user_id,
        'jti': token_id,
        'iat': now,
        'exp': now + timedelta(seconds=REFRESH_TOKEN_TTL),
        'type': 'refresh',
        'device_id': device_id,
        'iss': 'https://auth.quixzoom.com',
        'aud': 'quixzoom-platform',
    }
    return jwt.encode(payload, JWT_PRIVATE_KEY, algorithm=JWT_ALGORITHM)

def verify_token(token: str, token_type: str = 'access') -> Optional[Dict]:
    try:
        payload = jwt.decode(
            token,
            JWT_PUBLIC_KEY,
            algorithms=[JWT_ALGORITHM],
            audience='quixzoom-platform',
            issuer='https://auth.quixzoom.com'
        )
        if payload.get('type') != token_type:
            return None
        # Check if token is blacklisted
        jti = payload.get('jti')
        if jti and redis_client.get(f'blacklist:{jti}'):
            return None
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

def set_auth_cookies(response: Response, access_token: str, refresh_token: str):
    """Set cookies that work across all quixzoom subdomains"""
    response.set_cookie(
        key='qz_access_token',
        value=access_token,
        max_age=ACCESS_TOKEN_TTL,
        httponly=True,
        secure=COOKIE_SECURE,
        samesite=COOKIE_SAMESITE,
        domain=COOKIE_DOMAIN,
        path='/'
    )
    response.set_cookie(
        key='qz_refresh_token',
        value=refresh_token,
        max_age=REFRESH_TOKEN_TTL,
        httponly=True,
        secure=COOKIE_SECURE,
        samesite=COOKIE_SAMESITE,
        domain=COOKIE_DOMAIN,
        path='/'
    )

def clear_auth_cookies(response: Response):
    """Clear auth cookies from all quixzoom domains"""
    for cookie_name in ['qz_access_token', 'qz_refresh_token']:
        response.delete_cookie(
            key=cookie_name,
            domain=COOKIE_DOMAIN,
            path='/'
        )

# ─── Authentication Dependency ───────────────────────────────────────────

async def get_current_user(request: Request) -> Optional[Dict]:
    """Extract and verify user from cookie or Authorization header"""
    token = None
    
    # Try cookie first (for web)
    token = request.cookies.get('qz_access_token')
    
    # Try Authorization header (for API/app)
    if not token:
        auth_header = request.headers.get('Authorization', '')
        if auth_header.startswith('Bearer '):
            token = auth_header[7:]
    
    if not token:
        return None
    
    payload = verify_token(token, 'access')
    if not payload:
        return None
    
    # TODO: Fetch full user from database
    return {
        'id': payload['sub'],
        'email': payload['email'],
        'role': payload['role'],
    }

# ─── Routes ──────────────────────────────────────────────────────────────

@app.get('/health')
async def health():
    return {'status': 'ok', 'service': 'quixzoom-sso-auth', 'version': '2.0.0'}

@app.post('/auth/login')
async def login(request: LoginRequest, response: Response):
    """Login and set cross-domain cookies"""
    # TODO: Verify credentials against database
    # For now, mock implementation
    
    user_id = str(uuid.uuid4())
    token_id = generate_token_id()
    refresh_id = generate_token_id()
    
    access_token = create_access_token(user_id, request.email, 'zoomer', token_id)
    refresh_token = create_refresh_token(user_id, refresh_id, request.device_id)
    
    # Store refresh token in Redis
    redis_client.setex(
        f'refresh:{refresh_id}',
        REFRESH_TOKEN_TTL,
        user_id
    )
    
    # Set cookies for web clients
    set_auth_cookies(response, access_token, refresh_token)
    
    return TokenResponse(
        access_token=access_token,
        refresh_token=refresh_token,
        user={
            'id': user_id,
            'email': request.email,
            'first_name': 'Test',
            'last_name': 'User',
        }
    )

@app.post('/auth/refresh')
async def refresh_token(request: Request, response: Response):
    """Refresh access token using refresh token"""
    refresh_token = request.cookies.get('qz_refresh_token')
    
    if not refresh_token:
        auth_header = request.headers.get('Authorization', '')
        if auth_header.startswith('Bearer '):
            refresh_token = auth_header[7:]
    
    if not refresh_token:
        raise HTTPException(status_code=401, detail='No refresh token provided')
    
    payload = verify_token(refresh_token, 'refresh')
    if not payload:
        raise HTTPException(status_code=401, detail='Invalid refresh token')
    
    # Check if refresh token is in Redis
    jti = payload.get('jti')
    user_id = redis_client.get(f'refresh:{jti}')
    
    if not user_id:
        raise HTTPException(status_code=401, detail='Refresh token revoked')
    
    # Rotate refresh token (security best practice)
    new_token_id = generate_token_id()
    new_refresh_id = generate_token_id()
    
    # Delete old refresh token
    redis_client.delete(f'refresh:{jti}')
    
    # Create new tokens
    access_token = create_access_token(user_id, payload.get('email', ''), 'zoomer', new_token_id)
    new_refresh_token = create_refresh_token(user_id, new_refresh_id, payload.get('device_id'))
    
    # Store new refresh token
    redis_client.setex(f'refresh:{new_refresh_id}', REFRESH_TOKEN_TTL, user_id)
    
    # Update cookies
    set_auth_cookies(response, access_token, new_refresh_token)
    
    return TokenResponse(
        access_token=access_token,
        refresh_token=new_refresh_token,
        user={'id': user_id, 'email': payload.get('email', '')}
    )

@app.post('/auth/logout')
async def logout(request: Request, response: Response):
    """Logout and clear all cookies"""
    # Blacklist the access token
    access_token = request.cookies.get('qz_access_token')
    if access_token:
        payload = verify_token(access_token)
        if payload and payload.get('jti'):
            redis_client.setex(f'blacklist:{payload["jti"]}', ACCESS_TOKEN_TTL, '1')
    
    # Delete refresh token
    refresh_token = request.cookies.get('qz_refresh_token')
    if refresh_token:
        payload = verify_token(refresh_token, 'refresh')
        if payload and payload.get('jti'):
            redis_client.delete(f'refresh:{payload["jti"]}')
    
    clear_auth_cookies(response)
    
    return {'status': 'logged_out'}

@app.get('/auth/me')
async def get_me(current_user: Optional[Dict] = Depends(get_current_user)):
    """Get current user info"""
    if not current_user:
        raise HTTPException(status_code=401, detail='Not authenticated')
    
    return {
        'authenticated': True,
        'user': current_user
    }

@app.get('/auth/check')
async def auth_check(current_user: Optional[Dict] = Depends(get_current_user)):
    """Quick auth check - returns 200 if authenticated, 401 if not"""
    if not current_user:
        raise HTTPException(status_code=401, detail='Not authenticated')
    
    return {
        'authenticated': True,
        'user': current_user
    }

@app.get('/auth/sso/initiate')
async def sso_initiate(
    redirect_url: str = 'https://www.quixzoom.com/mina-sidor',
    client_id: str = 'quixzoom-web'
):
    """Initiate SSO login flow - redirect to login page"""
    # Store the redirect URL in Redis for after login
    session_id = secrets.token_urlsafe(32)
    redis_client.setex(f'sso:{session_id}', 300, redirect_url)
    
    login_url = f'https://auth.quixzoom.com/login?session={session_id}&redirect={redirect_url}'
    return RedirectResponse(url=login_url)

@app.get('/auth/sso/callback')
async def sso_callback(
    code: str,
    state: str,
    response: Response
):
    """SSO callback - exchange code for tokens and redirect"""
    # Verify the code and get redirect URL
    redirect_url = redis_client.get(f'sso:{state}')
    if not redirect_url:
        raise HTTPException(status_code=400, detail='Invalid or expired session')
    
    # TODO: Verify code with identity provider
    
    # Set cookies and redirect
    token_id = generate_token_id()
    refresh_id = generate_token_id()
    
    # Mock user - replace with actual lookup
    user_id = str(uuid.uuid4())
    access_token = create_access_token(user_id, 'user@quixzoom.com', 'zoomer', token_id)
    refresh_token = create_refresh_token(user_id, refresh_id)
    
    set_auth_cookies(response, access_token, refresh_token)
    
    return RedirectResponse(url=redirect_url)

# ─── Public Key Endpoint ─────────────────────────────────────────────────

@app.post('/auth/sync')
async def sync_tokens(
    request: Request,
    response: Response,
    current_user: Optional[Dict] = Depends(get_current_user)
):
    """Sync mobile app tokens with web cookies
    
    Called by mobile app after login to enable seamless web auth.
    Sets cookies on .quixzoom.com domain so user is auto-logged in on web.
    """
    if not current_user:
        raise HTTPException(status_code=401, detail='Not authenticated')
    
    body = await request.json()
    refresh_token = body.get('refresh_token')
    
    # Create new tokens for web session
    token_id = generate_token_id()
    refresh_id = generate_token_id()
    
    access_token = create_access_token(
        current_user['id'],
        current_user['email'],
        current_user.get('role', 'zoomer'),
        token_id
    )
    
    refresh_token_web = create_refresh_token(
        current_user['id'],
        refresh_id,
        device_id='web_sync'
    )
    
    # Store refresh token
    redis_client.setex(f'refresh:{refresh_id}', REFRESH_TOKEN_TTL, current_user['id'])
    
    # Set cookies for web
    set_auth_cookies(response, access_token, refresh_token_web)
    
    return {
        'status': 'synced',
        'message': 'Web cookies set - user is now logged in on all quixzoom sites'
    }


@app.get('/auth/.well-known/jwks.json')
async def jwks():
    """JWKS endpoint for token verification by other services"""
    # TODO: Return proper JWKS format
    return {
        'keys': [
            {
                'kty': 'RSA',
                'use': 'sig',
                'kid': 'quixzoom-2026',
                'alg': 'RS256',
                # TODO: Add actual public key components
            }
        ]
    }

# ─── Error Handlers ──────────────────────────────────────────────────────

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content={'error': exc.detail}
    )

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content={'error': 'Internal server error'}
    )

if __name__ == '__main__':
    port = int(os.getenv('PORT', 8080))
    host = os.getenv('HOST', '0.0.0.0')
    uvicorn.run(app, host=host, port=port)