from datetime import datetime, timedelta
import jwt
from typing import Dict, List

class JWTService:
    """JWT Service for Neo4j authentication
    
    TODO: Security Enhancements Needed
    - Implement token refresh mechanism
    - Add token revocation capability
    - Add token validation checks
    - Implement rate limiting
    - Add audit logging for token generation/usage
    - Consider reducing token expiry time and implementing refresh tokens
    """

    def __init__(self, secret_key: str, algorithm: str = "HS256"):
        self.secret_key = secret_key
        self.algorithm = algorithm

    def generate_neo4j_token(self, user_data: Dict) -> str:
        """Generate JWT token for Neo4j database access"""
        payload = {
            "sub": user_data["email"],
            "roles": self._get_neo4j_roles(user_data["user_type"]),
            "iss": "supabase",
            "aud": "neo4j",
            "iat": datetime.utcnow(),
            "exp": datetime.utcnow() + timedelta(hours=24)
        }
        
        if "school_uuid" in user_data:
            payload["worker_db_name"] = f"cc.institutes.{user_data['school_uuid']}"
        
        return jwt.encode(payload, self.secret_key, algorithm=self.algorithm)

    def _get_neo4j_roles(self, user_type: str) -> List[str]:
        """Map user types to Neo4j roles"""
        role_mapping = {
            "cc_admin": ["admin", "reader", "writer"],
            "developer": ["developer", "reader", "writer"],
            "email_teacher": ["teacher", "reader", "writer"],
            "email_student": ["student", "reader"]
        }
        return role_mapping.get(user_type, ["reader"])