from datetime import datetime, timedelta, timezone
from typing import Optional
import time
import jwt
from passlib.context import CryptContext
import pyotp
from fastapi import Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
import config

# Password handling
pwd_context = CryptContext(schemes=["pbkdf2_sha256"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/auth/login")

def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password):
    return pwd_context.hash(password)

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    expire = datetime.now(timezone.utc) + (expires_delta or timedelta(minutes=15))
    to_encode.update({"exp": expire, "type": "access"})
    return jwt.encode(to_encode, config.SECRET_KEY, algorithm=config.ALGORITHM)

def create_refresh_token(data: dict):
    to_encode = data.copy()
    expire = datetime.now(timezone.utc) + timedelta(minutes=config.REFRESH_TOKEN_EXPIRE_MINUTES)
    to_encode.update({"exp": expire, "type": "refresh", "tv": _load_token_version()})
    return jwt.encode(to_encode, config.SECRET_KEY, algorithm=config.ALGORITHM)

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        if payload.get("type") != "access":
            raise credentials_exception
        username: str = payload.get("sub")
        if username is None or username != config.ADMIN_USER:
            raise credentials_exception
    except jwt.PyJWTError:
        raise credentials_exception
    return username

async def get_current_user_ws(token: str):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
    )
    try:
        payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        if payload.get("type") != "access":
            raise credentials_exception
        username: str = payload.get("sub")
        if username is None or username != config.ADMIN_USER:
            raise credentials_exception
    except jwt.PyJWTError:
        raise credentials_exception
    return username

# TOTP Persistence
AUTH_FILE = config.VOLUME_ROOT + "/docker/nas-dashboard/auth.json"
import json, os
import base64
import hashlib
from cryptography.fernet import Fernet

def _get_fernet():
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    from cryptography.hazmat.primitives import hashes
    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=b"nas-dashboard-fernet", iterations=480000)
    key = base64.urlsafe_b64encode(kdf.derive(config.SECRET_KEY.encode()))
    return Fernet(key)

def _get_legacy_fernet():
    key = base64.urlsafe_b64encode(hashlib.sha256(config.SECRET_KEY.encode()).digest())
    return Fernet(key)

def _encrypt(plaintext: str) -> str:
    if not plaintext:
        return ""
    return _get_fernet().encrypt(plaintext.encode()).decode()

def _decrypt(ciphertext: str) -> str:
    if not ciphertext:
        return ""
    try:
        return _get_fernet().decrypt(ciphertext.encode()).decode()
    except Exception:
        pass
    # Fallback: try legacy SHA256-based key
    try:
        return _get_legacy_fernet().decrypt(ciphertext.encode()).decode()
    except Exception:
        return ciphertext  # fallback for unencrypted legacy values

def _load_auth_data():
    try:
        if os.path.exists(AUTH_FILE):
            with open(AUTH_FILE, "r") as f:
                return json.load(f)
    except Exception:
        pass
    return {}

def _save_auth_data(data):
    os.makedirs(os.path.dirname(AUTH_FILE), exist_ok=True)
    with open(AUTH_FILE, "w") as f:
        json.dump(data, f)

def load_totp_secret():
    encrypted = _load_auth_data().get("totp_secret", "")
    if not encrypted:
        return config.TOTP_SECRET
    return _decrypt(encrypted)

def save_totp_secret(secret: str):
    data = _load_auth_data()
    data["totp_secret"] = _encrypt(secret) if secret else ""
    _save_auth_data(data)

def load_password_hash():
    return _load_auth_data().get("password_hash", "") or config.ADMIN_PASSWORD_HASH

def save_password_hash(hashed: str):
    data = _load_auth_data()
    data["password_hash"] = hashed
    _save_auth_data(data)
        
def verify_totp(token: str, secret: str = None):
    if secret is None:
        secret = load_totp_secret()
    if not secret:
        return True # 2FA not enabled
    totp = pyotp.TOTP(secret)
    if not totp.verify(token):
        return False
    # Replay protection: reject reuse within same 30s window
    current_ts = int(time.time()) // 30
    data = _load_auth_data()
    if data.get("last_totp_ts") == current_ts:
        return False
    data["last_totp_ts"] = current_ts
    _save_auth_data(data)
    return True

# Passkey credentials
def load_passkey_credentials():
    return _load_auth_data().get("passkey_credentials", [])

def save_passkey_credential(cred):
    data = _load_auth_data()
    creds = data.get("passkey_credentials", [])
    creds.append(cred)
    data["passkey_credentials"] = creds
    _save_auth_data(data)

def clear_passkey_credentials():
    data = _load_auth_data()
    data["passkey_credentials"] = []
    _save_auth_data(data)

# Token version for refresh token rotation
def _load_token_version() -> int:
    return _load_auth_data().get("token_version", 0)

def increment_token_version() -> int:
    data = _load_auth_data()
    v = data.get("token_version", 0) + 1
    data["token_version"] = v
    _save_auth_data(data)
    return v

def verify_token_version(tv: int) -> bool:
    return tv == _load_token_version()