import os
import secrets
import jwt
from datetime import datetime, timedelta
from argon2 import PasswordHasher
from fastapi import Request, HTTPException
from cryptography.fernet import Fernet
from dotenv import load_dotenv
import os

load_dotenv()

ph = PasswordHasher()
JWT_SECRET = os.getenv("JWT_SECRET", "super-secret-wolf-key")
ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY")
ALGORITHM = "HS256"

def decrypt_token(encrypted_token: str) -> str:
    if not ENCRYPTION_KEY:
        return encrypted_token
    try:
        f = Fernet(ENCRYPTION_KEY.encode())
        return f.decrypt(encrypted_token.encode()).decode()
    except:
        return encrypted_token

def hash_password(password: str) -> str:
    return ph.hash(password)

def verify_password(hash: str, password: str) -> bool:
    try:
        return ph.verify(hash, password)
    except:
        return False


class SecurityLayer:
    def __init__(self):
        self.allowed_ips = os.getenv("ALLOWED_IPS", "127.0.0.1").split(",")
        self.admin_hash = os.getenv("ADMIN_PASSWORD_HASH")
        self.pending_2fa = {} # {code: timestamp}

    async def check_ip(self, request: Request):
        client_ip = request.client.host
        if client_ip not in self.allowed_ips and "127.0.0.1" not in self.allowed_ips:
             if os.getenv("ENFORCE_IP_WHITELIST", "false").lower() == "true":
                raise HTTPException(status_code=403, detail="IP NOT WHITELISTED")

    def validate_admin(self, password: str) -> bool:
        if not self.admin_hash:
            return False
        return verify_password(self.admin_hash, password)

    def generate_2fa(self):
        code = "".join([str(secrets.randbelow(10)) for _ in range(6)])
        return code

    def create_access_token(self, data: dict):
        to_encode = data.copy()
        expire = datetime.utcnow() + timedelta(hours=12)
        to_encode.update({"exp": expire})
        return jwt.encode(to_encode, JWT_SECRET, algorithm=ALGORITHM)

    def decode_token(self, token: str):
        try:
            return jwt.decode(token, JWT_SECRET, algorithms=[ALGORITHM])
        except:
            return None