"""Secure encrypted storage for OAuth tokens.

This module provides encrypted storage for OAuth tokens using Fernet (AES-128).
Tokens are stored in an encrypted file with proper permissions (0600).
"""

import json
import os
from pathlib import Path
from typing import Dict, Optional

from cryptography.fernet import Fernet, InvalidToken

from .oauth import OAuthToken


class TokenStore:
    """Secure encrypted storage for OAuth tokens.

    Tokens are encrypted using Fernet (AES-128) and stored in a file with
    0600 permissions (owner read/write only).
    """

    def __init__(self, storage_path: Optional[Path] = None):
        """Initialize token store.

        Args:
            storage_path: Directory for token storage (default: ~/.config/groknroll)
        """
        if storage_path is None:
            storage_path = Path.home() / ".config" / "groknroll"

        self.storage_path = storage_path
        self.token_file = storage_path / "mcp_tokens.enc"
        self.key_file = storage_path / "oauth_key"

        # Ensure storage directory exists with proper permissions
        self.storage_path.mkdir(parents=True, exist_ok=True)
        os.chmod(self.storage_path, 0o700)  # drwx------

        # Load or create encryption key
        self.encryption_key = self._load_or_create_key()
        self.cipher = Fernet(self.encryption_key)

    def save_token(self, server_name: str, token: OAuthToken) -> None:
        """Save OAuth token for a server.

        Args:
            server_name: Name of the MCP server
            token: OAuth token to save
        """
        # Load existing tokens
        tokens = self._load_all_tokens()

        # Add/update token
        tokens[server_name] = token.to_dict()

        # Save encrypted
        self._save_all_tokens(tokens)

    def load_token(self, server_name: str) -> Optional[OAuthToken]:
        """Load OAuth token for a server.

        Args:
            server_name: Name of the MCP server

        Returns:
            OAuth token if found, None otherwise
        """
        tokens = self._load_all_tokens()
        token_data = tokens.get(server_name)

        if token_data is None:
            return None

        return OAuthToken.from_stored_dict(token_data)

    def delete_token(self, server_name: str) -> bool:
        """Delete OAuth token for a server.

        Args:
            server_name: Name of the MCP server

        Returns:
            True if token was deleted, False if not found
        """
        tokens = self._load_all_tokens()

        if server_name not in tokens:
            return False

        del tokens[server_name]
        self._save_all_tokens(tokens)
        return True

    def list_servers(self) -> list[str]:
        """List all servers with stored tokens.

        Returns:
            List of server names
        """
        tokens = self._load_all_tokens()
        return list(tokens.keys())

    def clear_all(self) -> None:
        """Delete all stored tokens."""
        if self.token_file.exists():
            self.token_file.unlink()

    def _load_or_create_key(self) -> bytes:
        """Load existing encryption key or create new one.

        Returns:
            Encryption key bytes
        """
        if self.key_file.exists():
            with open(self.key_file, "rb") as f:
                key = f.read()
            return key

        # Generate new key
        key = Fernet.generate_key()

        # Save with restricted permissions
        with open(self.key_file, "wb") as f:
            f.write(key)
        os.chmod(self.key_file, 0o600)  # -rw-------

        return key

    def _load_all_tokens(self) -> Dict[str, dict]:
        """Load all tokens from encrypted storage.

        Returns:
            Dictionary of server_name -> token_data
        """
        if not self.token_file.exists():
            return {}

        try:
            with open(self.token_file, "rb") as f:
                encrypted_data = f.read()

            # Decrypt
            decrypted_data = self.cipher.decrypt(encrypted_data)

            # Parse JSON
            return json.loads(decrypted_data.decode("utf-8"))

        except (InvalidToken, json.JSONDecodeError, UnicodeDecodeError):
            # Corrupted data, return empty
            return {}

    def _save_all_tokens(self, tokens: Dict[str, dict]) -> None:
        """Save all tokens to encrypted storage.

        Args:
            tokens: Dictionary of server_name -> token_data
        """
        # Convert to JSON
        json_data = json.dumps(tokens, indent=2)
        json_bytes = json_data.encode("utf-8")

        # Encrypt
        encrypted_data = self.cipher.encrypt(json_bytes)

        # Save with restricted permissions
        with open(self.token_file, "wb") as f:
            f.write(encrypted_data)
        os.chmod(self.token_file, 0o600)  # -rw-------

    def rotate_key(self) -> None:
        """Rotate encryption key and re-encrypt all tokens.

        This should be called periodically for security.
        """
        # Load all tokens with current key
        tokens = self._load_all_tokens()

        # Generate new key
        new_key = Fernet.generate_key()

        # Save new key
        with open(self.key_file, "wb") as f:
            f.write(new_key)
        os.chmod(self.key_file, 0o600)

        # Update cipher
        self.encryption_key = new_key
        self.cipher = Fernet(self.encryption_key)

        # Re-encrypt tokens with new key
        self._save_all_tokens(tokens)

    def get_token_info(self, server_name: str) -> Optional[dict]:
        """Get metadata about a stored token without decrypting sensitive data.

        Args:
            server_name: Name of the MCP server

        Returns:
            Dictionary with token metadata (expires_at, scope, etc.)
        """
        token = self.load_token(server_name)
        if token is None:
            return None

        return {
            "server": server_name,
            "token_type": token.token_type,
            "expires_at": token.expires_at().isoformat(),
            "is_expired": token.is_expired(),
            "has_refresh_token": token.refresh_token is not None,
            "scope": token.scope,
        }
