"""Key storage and management for 0b1 agents."""

import json
from pathlib import Path
from typing import Optional

from .protocol import DEFAULT_KEY_PATH
from .crypto import generate_keypair, private_to_address, private_to_public


def ensure_key_dir() -> None:
    """Ensure ~/.ob1 directory exists."""
    DEFAULT_KEY_PATH.parent.mkdir(parents=True, exist_ok=True)


def save_key(
    private_key: str,
    path: Optional[Path] = None,
    passphrase: Optional[str] = None
) -> Path:
    """
    Save key to JSON file.
    
    Args:
        private_key: "0x" + 64 hex chars
        path: Save location (default: ~/.ob1/key.json)
        passphrase: Encryption passphrase (not implemented in MVP)
    
    Returns:
        Path: Where key was saved
    """
    if path is None:
        path = DEFAULT_KEY_PATH
    
    path = Path(path).expanduser()
    path.parent.mkdir(parents=True, exist_ok=True)
    
    address = private_to_address(private_key)
    public_key = private_to_public(private_key)
    
    data = {
        "version": 1,
        "address": address,
        "public_key": public_key,
        "private_key": private_key,
        "encrypted": passphrase is not None
    }
    
    # TODO: Implement actual encryption with passphrase
    if passphrase:
        raise NotImplementedError("Encrypted keystore not yet implemented")
    
    with open(path, "w") as f:
        json.dump(data, f, indent=2)
    
    # Set restrictive permissions
    path.chmod(0o600)
    
    return path


def load_key(
    path: Optional[Path] = None,
    passphrase: Optional[str] = None
) -> str:
    """
    Load key from file.
    
    Args:
        path: Key file path (default: ~/.ob1/key.json)
        passphrase: Decryption passphrase
    
    Returns:
        str: Private key hex
    
    Raises:
        FileNotFoundError: If key file doesn't exist
        ValueError: If key file is invalid
    """
    if path is None:
        path = DEFAULT_KEY_PATH
    
    path = Path(path).expanduser()
    
    if not path.exists():
        raise FileNotFoundError(f"Key file not found: {path}")
    
    with open(path) as f:
        data = json.load(f)
    
    if data.get("encrypted"):
        if not passphrase:
            raise ValueError("Key is encrypted, passphrase required")
        raise NotImplementedError("Encrypted keystore not yet implemented")
    
    private_key = data.get("private_key")
    if not private_key:
        raise ValueError("Invalid key file: missing private_key")
    
    return private_key


def import_wallet(
    private_key: str,
    path: Optional[Path] = None
) -> Path:
    """
    Import existing wallet and save to file.
    
    Args:
        private_key: Private key hex
        path: Save location (default: ~/.ob1/key.json)
    
    Returns:
        Path: Where key was saved
    """
    # Validate key by deriving address
    try:
        private_to_address(private_key)
    except Exception as e:
        raise ValueError(f"Invalid private key: {e}")
    
    return save_key(private_key, path)


def key_exists(path: Optional[Path] = None) -> bool:
    """Check if key file exists."""
    if path is None:
        path = DEFAULT_KEY_PATH
    return Path(path).expanduser().exists()
