import os
import json
import time
import hmac
import hashlib
import base64
import requests
import logging
from typing import Optional, Dict, Any, List
import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from .types import AuthenticatedRequest, Project
from .utils.redis_connection import get_redis

logger = logging.getLogger(__name__)


class ProjectTokenPayload:
    """Project token payload structure"""
    
    def __init__(self, data: Dict[str, Any]):
        self.project_uuid: str = data.get('project_uuid', '')
        self.enabled_services: List[str] = data.get('enabled_services', [])
        self.secret_version: int = data.get('secret_version', 0)
        self.token_id: str = data.get('token_id', '')
        self.exp: int = data.get('exp', 0)
        self.iat: int = data.get('iat', 0)
        self.iss: str = data.get('iss', '')


class ProjectAuthMiddleware:
    """JWT-based Project-level authentication middleware"""
    
    def __init__(self):
        self.mercury_base_url = os.getenv('MERCURY_BASE_URL', 'http://localhost:3000')
        self.signature_secret = os.getenv('SIGNATURE_SHARED_SECRET', '')
        self.service_id = os.getenv('SERVICE_ID', '')
        self.jwks_cache_key = 'project_jwks_cache'
        self.jwks_cache_ttl = int(os.getenv('JWKS_CACHE_TTL', '18000'))  # 5 hours default
    
    async def authenticate(self, req: AuthenticatedRequest) -> None:
        """
        Authenticate project using x-project-token JWT header
        
        Args:
            req: AuthenticatedRequest object
            
        Raises:
            Exception: If authentication fails
        """
        try:
            logger.info("🚀 Starting project authentication...")
            
            # Extract token from x-project-token header
            headers = req.headers
            logger.info(f"📋 Available headers: {list(headers.keys())}")
            
            auth_header = headers.get('x-project-token')
            logger.info(f"🔑 x-project-token header present: {bool(auth_header)}")
            
            if not auth_header:
                logger.error("❌ No project token provided")
                raise Exception("No project token provided, required_header: 'x-project-token'")
            
            # Handle Bearer prefix
            token = auth_header[7:] if auth_header.startswith('Bearer ') else auth_header
            logger.info(f"🎫 Token extracted (length: {len(token)}, starts with Bearer: {auth_header.startswith('Bearer ')})")
            logger.info(f"🎫 Token preview: {token[:50]}...")
            
            if not token:
                logger.error("❌ Empty project token after processing")
                raise Exception('Empty project token')
            
            # Validate project token using cached JWKS
            logger.info("🔍 Starting token validation...")
            validation = await self._validate_project_token(token)
            
            logger.info(f"✅ Token validation result: {validation['is_valid']}")
            
            if not validation['is_valid']:
                logger.error(f"❌ Token validation failed: {validation['error']}")
                raise Exception(f"Invalid project token, message: {validation['error']}")
            
            # Check if current service is enabled for this project
            if not self.service_id:
                logger.error("❌ Service ID not configured")
                raise Exception('Service ID not configured')
            
            payload = validation['payload']
            logger.info(f"🎯 Checking service access - Service ID: {self.service_id}")
            logger.info(f"🎯 Enabled services: {payload.enabled_services}")
            
            has_access = self.service_id in payload.enabled_services
            logger.info(f"🎯 Service access granted: {has_access}")
            
            if not has_access:
                logger.error(f"❌ Service access denied")
                raise Exception(f"""
                    error: Service access denied,
                    service_id: {self.service_id},
                    project: {payload.project_uuid},
                    enabled_services: {payload.enabled_services}
                """)
            
            # Inject project context into request (NO user context)
            logger.info("📝 Creating project context...")
            req.project = Project(
                project_uuid=payload.project_uuid,
                enabled_services=payload.enabled_services,
                secret_version=payload.secret_version,
                token_id=payload.token_id,
                expires_at=payload.exp
            )
            
            logger.info(f"✅ Authentication successful for project: {payload.project_uuid}")
            
        except Exception as e:
            logger.error(f"❌ Project authentication failed: {str(e)}", exc_info=True)
            raise Exception(f"Authentication service error: {str(e)}")
    
    async def _validate_project_token(self, token: str) -> Dict[str, Any]:
        """
        Validate project token using cached JWKS + RSA verification
        
        Args:
            token: JWT token string
            
        Returns:
            Dict with validation results
        """
        try:
            # Get public key from cached JWKS
            public_key = await self._get_public_key_from_cache(token)
            
            # Verify JWT with RSA public key
            try:
                verified = jwt.decode(
                    token,
                    public_key,
                    algorithms=["RS512"],
                    options={
                        "verify_exp": True,
                        "verify_aud": False,  # Skip audience verification to match Node.js behavior
                        "verify_iss": False   # Skip issuer verification to match Node.js behavior
                    }
                )
            except jwt.ExpiredSignatureError:
                return {'is_valid': False, 'error': 'Token has expired'}
            except jwt.InvalidTokenError as e:
                return {'is_valid': False, 'error': f'Invalid token: {str(e)}'}
            
            payload = ProjectTokenPayload(verified)
            
            # Validate project token structure
            if not payload.project_uuid or not payload.token_id or not payload.enabled_services:
                return {'is_valid': False, 'error': 'Invalid project token structure'}
            
            # Check secret version
            current_secret_version = await self._get_current_project_secret_version(payload.project_uuid)
            if current_secret_version > 0 and payload.secret_version < current_secret_version:
                return {
                    'is_valid': False,
                    'error': f'Token secret version outdated (token: {payload.secret_version}, current: {current_secret_version}) - re-authentication required'
                }
            
            # Check if token is in Redis cache (for revocation)
            redis = await get_redis()
            token_key = f"project_token:{payload.token_id}"
            
            logger.info(f"🔍 Checking token existence in Redis - Key: {token_key}")
            
            # Use the Redis client's exists method, but handle it properly
            try:
                token_exists = redis.client.exists(token_key)
                logger.info(f"🔍 Redis EXISTS result: {token_exists} (type: {type(token_exists)})")
            except Exception as e:
                logger.error(f"❌ Redis EXISTS error: {str(e)}")
                # Fallback: try to get the key directly
                try:
                    fallback_result = redis.client.get(token_key)
                    token_exists = 1 if fallback_result is not None else 0
                    logger.info(f"🔍 Fallback GET result: {fallback_result is not None} -> exists: {token_exists}")
                except Exception as fallback_error:
                    logger.error(f"❌ Fallback GET error: {str(fallback_error)}")
                    # If Redis is completely failing, allow the token for now
                    logger.warning("⚠️ Redis check failed, allowing token")
                    token_exists = 1
            
            if token_exists == 0:
                logger.error(f"❌ Token not found in Redis: {token_key}")
                return {'is_valid': False, 'error': 'Token has been revoked'}
            
            logger.info(f"✅ Token exists in Redis: {token_key}")
            
            return {'is_valid': True, 'payload': payload}
            
        except Exception as e:
            error_msg = str(e) if isinstance(e, Exception) else 'Token validation failed'
            return {'is_valid': False, 'error': error_msg}
    
    async def _get_public_key_from_cache(self, token: str) -> str:
        """
        Get RSA public key from cached JWKS (with 5+ hour caching)
        
        Args:
            token: JWT token to extract kid from
            
        Returns:
            PEM formatted public key
        """
        try:
            # Extract kid from JWT header
            header = self._decode_jwt_header(token)
            kid = header.get('kid')
            
            logger.info(f"🔑 Extracted kid from JWT header: {kid}")
            
            if not kid:
                raise Exception('Missing key ID in token header')
            
            # Check Redis for cached JWKS
            redis = await get_redis()
            cached_jwks = await redis.get(self.jwks_cache_key)
            
            logger.info(f"💾 Cached JWKS type: {type(cached_jwks)}")
            logger.info(f"💾 Cached JWKS present: {bool(cached_jwks)}")
            
            if cached_jwks:
                # The RedisConnectionManager.get() already handles JSON parsing
                if isinstance(cached_jwks, dict):
                    jwks_data = cached_jwks
                    logger.info("💾 Using cached JWKS as dict")
                elif isinstance(cached_jwks, str):
                    # It's a string, parse it
                    jwks_data = json.loads(cached_jwks)
                    logger.info("💾 Parsed cached JWKS from string")
                else:
                    jwks_data = cached_jwks  # Assume it's already parsed
                    logger.info("💾 Using cached JWKS as-is")
                key_store = jwks_data
            else:
                # Fetch fresh JWKS from Mercury and cache it
                logger.info("🌐 No cached JWKS, fetching fresh")
                key_store = await self._fetch_and_cache_jwks()
            
            logger.info(f"🔍 Available keys in JWKS: {[k.get('kid', 'no-kid') for k in key_store.get('keys', [])]}")
            
            # Find the specific key
            key = None
            for jwk in key_store.get('keys', []):
                if jwk.get('kid') == kid:
                    key = jwk
                    break
            
            if not key:
                raise Exception(f"Key {kid} not found in JWKS. Available keys: {[k.get('kid', 'no-kid') for k in key_store.get('keys', [])]}")
            
            logger.info(f"🔑 Found JWK: {key.keys()}")
            
            # Convert JWK to PEM format
            public_key_pem = self._jwk_to_pem(key)
            logger.info(f"🔐 Converted JWK to PEM (length: {len(public_key_pem)})")
            return public_key_pem
            
        except Exception as e:
            logger.error(f"❌ Error in _get_public_key_from_cache: {str(e)}", exc_info=True)
            raise Exception(f"Failed to get public key: {str(e)}")
    
    async def _fetch_and_cache_jwks(self) -> Dict[str, Any]:
        """
        Fetch JWKS from Mercury and cache in Redis
        
        Returns:
            JWKS data
        """
        try:
            logger.info("🌐 Fetching JWKS from Mercury service...")
            
            path = 'auth/project/.well-known/jwks.json'
            jwks_uri = f"{self.mercury_base_url}/{path}"
            timestamp = str(int(time.time() * 1000))
            signature_input = f'GET/{path}{timestamp}'
            
            logger.info(f"🔗 JWKS URI: {jwks_uri}")
            logger.info(f"⏰ Timestamp: {timestamp}")
            logger.info(f"📝 Signature input: {signature_input}")
            
            signature = hmac.new(
                self.signature_secret.encode(),
                signature_input.encode(),
                hashlib.sha256
            ).hexdigest()
            
            logger.info(f"🔐 Generated signature: {signature[:16]}...")
            
            headers = {
                'Accept': 'application/json',
                'User-Agent': 'Python-JWT-Strategy/1.0',
                'X-Timestamp': timestamp,
                'X-Signature': signature,
            }
            
            logger.info(f"📤 Request headers: {list(headers.keys())}")
            
            response = requests.get(jwks_uri, headers=headers, timeout=10)
            
            logger.info(f"📥 Response status: {response.status_code}")
            logger.info(f"📥 Response headers: {dict(response.headers)}")
            
            if response.status_code != 200:
                logger.error(f"❌ HTTP error {response.status_code}: {response.text}")
                if response.status_code == 0:  # Connection refused
                    raise Exception('Mercury service unavailable')
                raise Exception(f"HTTP {response.status_code}: {response.text}")
            
            logger.info(f"📄 Response text (first 500 chars): {response.text[:500]}")
            
            jwks_data = response.json()
            
            logger.info(f"✅ JWKS JSON parsed. Type: {type(jwks_data)}")
            logger.info(f"📋 JWKS keys: {list(jwks_data.keys()) if isinstance(jwks_data, dict) else 'Not a dict'}")
            
            if not jwks_data or 'keys' not in jwks_data:
                logger.error(f"❌ Invalid JWKS response structure: {jwks_data}")
                raise Exception('Invalid JWKS response')
            
            logger.info(f"🔑 Number of keys in JWKS: {len(jwks_data['keys'])}")
            
            # Cache JWKS in Redis for 5+ hours
            logger.info("💾 Caching JWKS in Redis...")
            redis = await get_redis()
            
            # Use the RedisConnectionManager.set method with EX option for TTL
            await redis.set(
                self.jwks_cache_key,
                jwks_data,  # Pass the dict directly, it will be JSON stringified
                {'EX': self.jwks_cache_ttl}
            )
            
            logger.info(f"✅ JWKS cached with TTL: {self.jwks_cache_ttl} seconds")
            
            # Verify what was cached
            verify_cached = await redis.get(self.jwks_cache_key)
            logger.info(f"🔍 Verification - cached type: {type(verify_cached)}")
            
            return jwks_data
            
        except requests.exceptions.ConnectionError as e:
            logger.error(f"❌ Connection error: {str(e)}")
            raise Exception('Mercury service unavailable')
        except requests.exceptions.Timeout as e:
            logger.error(f"❌ Timeout error: {str(e)}")
            raise Exception('Mercury service timeout')
        except requests.exceptions.RequestException as e:
            logger.error(f"❌ Request error: {str(e)}")
            raise Exception(f"Failed to fetch JWKS: {str(e)}")
        except json.JSONDecodeError as e:
            logger.error(f"❌ JSON decode error: {str(e)}")
            logger.error(f"Response text: {response.text}")
            raise Exception(f"Invalid JSON response from Mercury: {str(e)}")
        except Exception as e:
            logger.error(f"❌ Unexpected error: {str(e)}", exc_info=True)
            raise Exception(f"Failed to fetch JWKS: {str(e)}")
    
    def _decode_jwt_header(self, token: str) -> Dict[str, Any]:
        """
        Decode JWT header to extract kid
        
        Args:
            token: JWT token string
            
        Returns:
            Decoded header dict
        """
        parts = token.split('.')
        if len(parts) != 3:
            raise Exception('Invalid JWT format')
        
        # Add padding if needed
        header_b64 = parts[0]
        padding = len(header_b64) % 4
        if padding:
            header_b64 += '=' * (4 - padding)
        
        header_json = base64.b64decode(header_b64).decode()
        return json.loads(header_json)
    
    def _jwk_to_pem(self, jwk: Dict[str, Any]) -> str:
        """
        Convert JWK to PEM format public key
        
        Args:
            jwk: JSON Web Key dict
            
        Returns:
            PEM formatted public key
        """
        try:
            logger.info(f"🔧 Converting JWK to PEM. JWK keys: {list(jwk.keys())}")
            
            # Extract RSA components from JWK
            n_b64 = jwk['n']
            e_b64 = jwk['e']
            
            logger.info(f"🔧 n_b64 length: {len(n_b64)}, e_b64 length: {len(e_b64)}")
            
            # Add padding if needed
            n_b64 += '=' * (4 - len(n_b64) % 4) if len(n_b64) % 4 else ''
            e_b64 += '=' * (4 - len(e_b64) % 4) if len(e_b64) % 4 else ''
            
            logger.info(f"🔧 After padding - n_b64 length: {len(n_b64)}, e_b64 length: {len(e_b64)}")
            
            n_bytes = base64.urlsafe_b64decode(n_b64)
            e_bytes = base64.urlsafe_b64decode(e_b64)
            
            logger.info(f"🔧 Decoded bytes - n: {len(n_bytes)} bytes, e: {len(e_bytes)} bytes")
            
            # Convert to integers
            n = int.from_bytes(n_bytes, 'big')
            e = int.from_bytes(e_bytes, 'big')
            
            logger.info(f"🔧 RSA components - n: {n}, e: {e}")
            
            # Create RSA public key
            public_numbers = rsa.RSAPublicNumbers(e, n)
            logger.info(f"🔧 Created RSAPublicNumbers: {type(public_numbers)}")
            
            public_key = public_numbers.public_key()
            logger.info(f"🔧 Created public key: {type(public_key)}")
            logger.info(f"🔧 Public key attributes: {[attr for attr in dir(public_key) if not attr.startswith('_')]}")
            
            # Try different method names for serialization (API changed in newer versions)
            pem = None
            if hasattr(public_key, 'public_bytes'):
                logger.info("🔧 Using public_bytes method")
                pem = public_key.public_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo
                )
            elif hasattr(public_key, 'public_key_bytes'):
                logger.info("🔧 Using public_key_bytes method")
                pem = public_key.public_key_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo
                )
            else:
                raise Exception(f"Neither public_bytes nor public_key_bytes method found. Available methods: {[attr for attr in dir(public_key) if 'public' in attr.lower()]}")
            
            result = pem.decode('utf-8')
            logger.info(f"🔧 Successfully converted to PEM (length: {len(result)})")
            logger.info(f"🔧 PEM preview: {result[:100]}...")
            
            return result
            
        except Exception as e:
            logger.error(f"❌ Error in _jwk_to_pem: {str(e)}", exc_info=True)
            raise Exception(f"Failed to convert JWK to PEM: {str(e)}")
    
    async def _get_current_project_secret_version(self, project_uuid: str) -> int:
        """
        Get current project secret version from Redis (cached by Mercury)
        
        Args:
            project_uuid: Project UUID
            
        Returns:
            Current secret version
        """
        try:
            redis = await get_redis()
            cache_key = f"project_secret_version:{project_uuid}"
            
            cached_version = await redis.get(cache_key)
            
            if cached_version is not None:
                # The RedisConnectionManager.get() might return int or string
                if isinstance(cached_version, int):
                    return cached_version
                elif isinstance(cached_version, str):
                    return int(cached_version)
                else:
                    # Try to convert whatever it is
                    return int(str(cached_version))
            
            # If not in Redis, this means Mercury hasn't cached it yet
            # This shouldn't happen in normal flow, but fallback to allowing the token
            # Mercury's Kafka handler will eventually populate this
            return 0  # Default to allow if version not found
            
        except Exception as e:
            raise Exception(f"Failed to get project secret version: {str(e)}")
    
    async def refresh_jwks_cache(self) -> None:
        """Manually refresh JWKS cache"""
        await self._fetch_and_cache_jwks()
    
    def set_cache_ttl(self, seconds: int) -> None:
        """Update JWKS cache TTL (can be increased beyond 5 hours)"""
        self.jwks_cache_ttl = seconds