"""Storage layer for wallet persistence."""

import json
import os
import re
import shutil
from dataclasses import asdict, dataclass, field
from datetime import datetime, UTC
from pathlib import Path
from typing import Optional

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64


DEFAULT_DIR = Path.home() / ".aqua-mcp"
SWAP_ID_PATTERN = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')


@dataclass
class WalletData:
    """Wallet data structure."""
    name: str
    network: str  # "mainnet" or "testnet"
    descriptor: str  # CT descriptor (Liquid)
    btc_descriptor: Optional[str] = None  # BIP84 external descriptor (Bitcoin)
    btc_change_descriptor: Optional[str] = None  # BIP84 change descriptor (Bitcoin)
    encrypted_mnemonic: Optional[str] = None  # Encrypted, if full wallet
    watch_only: bool = False
    created_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "WalletData":
        # Backward compatibility: old wallet files may not have btc_* fields
        data = {**data}
        data.setdefault("btc_descriptor", None)
        data.setdefault("btc_change_descriptor", None)
        return cls(**data)


@dataclass
class Config:
    """Global configuration."""
    network: str = "mainnet"
    default_wallet: str = "default"
    electrum_url: Optional[str] = None
    auto_sync: bool = True

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "Config":
        return cls(**data)


def _validate_wallet_name(name: str) -> str:
    """Validate wallet name to prevent path traversal."""
    if not re.fullmatch(r'[a-zA-Z0-9_-]{1,64}', name):
        raise ValueError(
            f"Invalid wallet name '{name}'. "
            "Use only letters, numbers, hyphens and underscores (max 64 chars)."
        )
    return name


class Storage:
    """Handles wallet and config persistence."""

    def __init__(self, base_dir: Optional[Path] = None):
        self.base_dir = base_dir or DEFAULT_DIR
        self.wallets_dir = self.base_dir / "wallets"
        self.cache_dir = self.base_dir / "cache"
        self.swaps_dir = self.base_dir / "swaps"
        self.ankara_swaps_dir = self.base_dir / "ankara_swaps"
        self.lightning_swaps_dir = self.base_dir / "lightning_swaps"
        self.config_path = self.base_dir / "config.json"
        self._ensure_dirs()

    def _ensure_dirs(self):
        """Create necessary directories with restricted permissions."""
        self.base_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
        os.chmod(self.base_dir, 0o700)
        self.wallets_dir.mkdir(exist_ok=True, mode=0o700)
        os.chmod(self.wallets_dir, 0o700)
        self.cache_dir.mkdir(exist_ok=True, mode=0o700)
        os.chmod(self.cache_dir, 0o700)
        self.swaps_dir.mkdir(exist_ok=True, mode=0o700)
        os.chmod(self.swaps_dir, 0o700)
        self.ankara_swaps_dir.mkdir(exist_ok=True, mode=0o700)
        os.chmod(self.ankara_swaps_dir, 0o700)
        self.lightning_swaps_dir.mkdir(exist_ok=True, mode=0o700)
        os.chmod(self.lightning_swaps_dir, 0o700)

    def _derive_key(self, passphrase: str, salt: bytes) -> bytes:
        """Derive encryption key from passphrase."""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=480000,
        )
        return base64.urlsafe_b64encode(kdf.derive(passphrase.encode()))

    def encrypt_mnemonic(self, mnemonic: str, passphrase: str) -> str:
        """Encrypt mnemonic with passphrase."""
        salt = os.urandom(16)
        key = self._derive_key(passphrase, salt)
        f = Fernet(key)
        encrypted = f.encrypt(mnemonic.encode())
        # Store salt + encrypted data
        return base64.b64encode(salt + encrypted).decode()

    def decrypt_mnemonic(self, encrypted: str, passphrase: str) -> str:
        """Decrypt mnemonic with passphrase."""
        data = base64.b64decode(encrypted)
        salt = data[:16]
        encrypted_data = data[16:]
        key = self._derive_key(passphrase, salt)
        f = Fernet(key)
        return f.decrypt(encrypted_data).decode()

    def store_mnemonic(self, mnemonic: str, passphrase: Optional[str] = None) -> str:
        """Store mnemonic, encrypting only when a passphrase is provided."""
        if passphrase:
            return self.encrypt_mnemonic(mnemonic, passphrase)
        return "plain:" + base64.b64encode(mnemonic.encode()).decode()

    def retrieve_mnemonic(self, stored: str, passphrase: Optional[str] = None) -> str:
        """Retrieve mnemonic stored by store_mnemonic."""
        if stored.startswith("plain:"):
            return base64.b64decode(stored[6:]).decode()
        if not passphrase:
            raise ValueError("Passphrase required to decrypt mnemonic")
        return self.decrypt_mnemonic(stored, passphrase)

    def is_mnemonic_encrypted(self, stored: str) -> bool:
        """Check whether a stored mnemonic requires a passphrase."""
        return not stored.startswith("plain:")

    # Config operations

    def load_config(self) -> Config:
        """Load global configuration."""
        if self.config_path.exists():
            with open(self.config_path) as f:
                return Config.from_dict(json.load(f))
        return Config()

    def _atomic_write_json(self, path: Path, data: dict) -> None:
        """Atomically write JSON data to a file with restricted permissions."""
        tmp_path = path.with_suffix(path.suffix + ".tmp")
        try:
            with open(tmp_path, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)
                f.flush()
                os.fsync(f.fileno())
            if hasattr(os, "chmod"):
                try:
                    os.chmod(tmp_path, 0o600)
                except OSError:
                    pass
            os.replace(tmp_path, path)
            if hasattr(os, "chmod"):
                try:
                    os.chmod(path, 0o600)
                except OSError:
                    pass
        except Exception:
            if tmp_path.exists():
                try:
                    tmp_path.unlink()
                except OSError:
                    pass
            raise

    def save_config(self, config: Config):
        self._atomic_write_json(self.config_path, config.to_dict())

    # Wallet operations

    def _wallet_path(self, name: str) -> Path:
        """Get path to wallet file."""
        _validate_wallet_name(name)
        return self.wallets_dir / f"{name}.json"

    def wallet_exists(self, name: str) -> bool:
        """Check if wallet exists."""
        return self._wallet_path(name).exists()

    def list_wallets(self) -> list[str]:
        """List all wallet names."""
        return [
            p.stem for p in self.wallets_dir.glob("*.json")
            if re.fullmatch(r'[a-zA-Z0-9_-]{1,64}', p.stem)
        ]

    def load_wallet(self, name: str) -> Optional[WalletData]:
        """Load wallet data."""
        path = self._wallet_path(name)
        if not path.exists():
            return None
        with open(path) as f:
            return WalletData.from_dict(json.load(f))

    def save_wallet(self, wallet: WalletData):
        path = self._wallet_path(wallet.name)
        self._atomic_write_json(path, wallet.to_dict())

    def delete_wallet(self, name: str) -> bool:
        """Delete wallet and its cache directory."""
        path = self._wallet_path(name)
        if not path.exists():
            return False
        path.unlink()
        cache_path = self.cache_dir / name
        if cache_path.is_dir():
            shutil.rmtree(cache_path)
        return True

    # Swap operations

    def _swap_path(self, swap_id: str) -> Path:
        """Get path to swap file, validating the ID to prevent path traversal."""
        if not SWAP_ID_PATTERN.fullmatch(swap_id):
            raise ValueError(
                f"Invalid swap ID '{swap_id}'. "
                "Use only letters, numbers, hyphens and underscores (max 128 chars)."
            )
        return self.swaps_dir / f"{swap_id}.json"

    def save_swap(self, swap) -> None:
        """Save swap data for recovery."""
        path = self._swap_path(swap.swap_id)
        self._atomic_write_json(path, swap.to_dict())

    def load_swap(self, swap_id: str):
        """Load swap data. Returns SwapInfo or None."""
        from .boltz import SwapInfo

        path = self._swap_path(swap_id)
        if not path.exists():
            return None
        with open(path) as f:
            return SwapInfo(**json.load(f))

    def list_swaps(self) -> list[str]:
        """List all swap IDs."""
        return [
            p.stem for p in self.swaps_dir.glob("*.json")
            if SWAP_ID_PATTERN.fullmatch(p.stem)
        ]

    # Ankara swap operations

    def _ankara_swap_path(self, swap_id: str) -> Path:
        """Get path to Ankara swap file, validating the ID to prevent path traversal."""
        if not SWAP_ID_PATTERN.fullmatch(swap_id):
            raise ValueError(
                f"Invalid swap ID '{swap_id}'. "
                "Use only letters, numbers, hyphens and underscores (max 128 chars)."
            )
        return self.ankara_swaps_dir / f"{swap_id}.json"

    def save_ankara_swap(self, swap) -> None:
        """Save Ankara swap data for recovery."""
        path = self._ankara_swap_path(swap.swap_id)
        self._atomic_write_json(path, swap.to_dict())

    def load_ankara_swap(self, swap_id: str):
        """Load Ankara swap data. Returns AnkaraSwapInfo or None."""
        from .ankara import AnkaraSwapInfo

        path = self._ankara_swap_path(swap_id)
        if not path.exists():
            return None
        with open(path) as f:
            return AnkaraSwapInfo(**json.load(f))

    def list_ankara_swaps(self) -> list[str]:
        """List all Ankara swap IDs."""
        return [
            p.stem for p in self.ankara_swaps_dir.glob("*.json")
            if SWAP_ID_PATTERN.fullmatch(p.stem)
        ]

    # Lightning swap operations

    def _lightning_swap_path(self, swap_id: str) -> Path:
        """Get path to Lightning swap file, validating the ID to prevent path traversal."""
        if not SWAP_ID_PATTERN.fullmatch(swap_id):
            raise ValueError(
                f"Invalid swap ID '{swap_id}'. "
                "Use only letters, numbers, hyphens and underscores (max 128 chars)."
            )
        return self.lightning_swaps_dir / f"{swap_id}.json"

    def save_lightning_swap(self, swap) -> None:
        """Save Lightning swap data for recovery."""
        path = self._lightning_swap_path(swap.swap_id)
        self._atomic_write_json(path, swap.to_dict())

    def load_lightning_swap(self, swap_id: str):
        """Load Lightning swap data. Returns LightningSwap or None."""
        from .lightning import LightningSwap

        path = self._lightning_swap_path(swap_id)
        if not path.exists():
            return None
        with open(path) as f:
            return LightningSwap.from_dict(json.load(f))

    def list_lightning_swaps(self) -> list[str]:
        """List all Lightning swap IDs."""
        return [
            p.stem for p in self.lightning_swaps_dir.glob("*.json")
            if SWAP_ID_PATTERN.fullmatch(p.stem)
        ]

    # Cache operations

    def get_cache_path(self, wallet_name: str) -> Path:
        """Get cache directory for wallet."""
        _validate_wallet_name(wallet_name)
        cache_path = self.cache_dir / wallet_name
        cache_path.mkdir(exist_ok=True, mode=0o700)
        return cache_path
