#!/usr/bin/env python3
"""
Rook Vault Setup — Full Security Onboarding

Orchestrates the complete security setup:
1. Generate BIP39 mnemonic
2. Optional key file
3. Authenticate Google Drive, get folder ID
4. Derive all keys (Bitwarden password, LUKS passphrase, backup key)
5. Create LUKS container
6. Start Vaultwarden inside container
7. Initialize vault with derived password
8. Import existing secrets from ~/.hermes/secrets/
9. Seal LUKS passphrase with fTPM (or PIN fallback)
10. Run first backup

This is the "Full Security" path from first-boot.sh.
"""

import json
import logging
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional

# Add parent paths
sys.path.insert(0, str(Path(__file__).parent))

from keygen import RookKeyDerivation, generate_mnemonic, hash_keyfile
from luks_manager import LuksContainer

logging.basicConfig(level=logging.INFO, format="[vault] %(message)s")
logger = logging.getLogger(__name__)

HERMES_HOME = Path.home() / ".hermes"
SECRETS_DIR = HERMES_HOME / "secrets"
VAULT_STATE_FILE = HERMES_HOME / ".vault-state.json"
COMPOSE_FILE = Path(__file__).parent / "docker-compose.vaultwarden.yml"


def save_vault_state(state: dict) -> None:
    """Save vault state (non-sensitive metadata only)."""
    VAULT_STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(VAULT_STATE_FILE, "w") as f:
        json.dump(state, f, indent=2)
    os.chmod(VAULT_STATE_FILE, 0o600)


def load_vault_state() -> dict:
    """Load vault state."""
    if VAULT_STATE_FILE.exists():
        with open(VAULT_STATE_FILE) as f:
            return json.load(f)  # type: ignore[no-any-return]
    return {}


def step_generate_mnemonic(word_count: int = 12) -> str:
    """Generate and display mnemonic for user to write down."""
    strength = 128 if word_count == 12 else 256
    mnemonic = generate_mnemonic(strength)

    print("\n" + "=" * 50)
    print("  YOUR RECOVERY WORDS")
    print("=" * 50)
    print()
    words = mnemonic.split()
    for i, word in enumerate(words, 1):
        print(f"    {i:2}. {word}")
    print()
    print("  WRITE THESE DOWN ON PAPER.")
    print("  Do NOT photograph or store digitally.")
    print("  You need these to recover your agent.")
    print("=" * 50)
    print()

    return mnemonic


def step_get_keyfile() -> Optional[str]:
    """Ask user for optional key file."""
    print("You can add a key file for extra security.")
    print("This can be any file — a photo, a poem, anything personal.")
    print()
    choice = input("Add a key file? [y/N]: ").strip().lower()
    if choice != "y":
        return None

    path = input("Path to key file: ").strip()
    if not Path(path).exists():
        print(f"File not found: {path}")
        return None

    # Show hash for verification
    file_hash = hash_keyfile(path)
    print(f"Key file accepted. SHA256: {file_hash.hex()[:16]}...")
    return path


def step_get_folder_id() -> str:
    """Get Google Drive folder ID (from rclone or create new)."""
    print("\nChecking Google Drive connection...")

    # Try to find existing rook-backups folder
    try:
        result = subprocess.run(
            ["rclone", "lsjson", "gdrive:rook-backups/", "--max-depth", "0"],
            capture_output=True,
            text=True,
            timeout=15,
        )
        if result.returncode == 0:
            # Get the folder ID
            result2 = subprocess.run(
                ["rclone", "lsjson", "gdrive:", "--max-depth", "1", "--dirs-only"],
                capture_output=True,
                text=True,
                timeout=15,
            )
            if result2.returncode == 0:
                folders = json.loads(result2.stdout)
                for f in folders:
                    if f.get("Name") == "rook-backups":
                        folder_id = f.get("ID", "")
                        if folder_id:
                            print(f"Found Google Drive folder: rook-backups (ID: {folder_id[:12]}...)")
                            return folder_id  # type: ignore[no-any-return]
    except (subprocess.TimeoutExpired, json.JSONDecodeError, FileNotFoundError):
        pass

    # Fallback: use a hash of the email/hostname as folder ID
    import hashlib

    hostname = os.uname().nodename
    user = os.environ.get("USER", "unknown")
    fallback_id = hashlib.sha256(f"{user}@{hostname}".encode()).hexdigest()[:24]
    print(f"Using device-derived folder ID: {fallback_id[:12]}...")
    return fallback_id


def step_create_vault(luks_passphrase: str, size_mb: int = 500) -> LuksContainer:
    """Create the LUKS container."""
    container = LuksContainer()

    if container.exists:
        print("LUKS container already exists. Skipping creation.")
        return container

    print(f"\nCreating {size_mb}MB encrypted vault...")
    container.create(luks_passphrase, size_mb)
    print("Vault created successfully.")
    return container


def step_start_vaultwarden(admin_token: str) -> bool:
    """Start Vaultwarden inside the LUKS container."""
    print("\nStarting Vaultwarden...")

    env = os.environ.copy()
    env["VAULTWARDEN_ADMIN_TOKEN"] = admin_token

    try:
        subprocess.run(
            ["docker", "compose", "-f", str(COMPOSE_FILE), "up", "-d"],
            env=env,
            check=True,
            capture_output=True,
            text=True,
        )
    except subprocess.CalledProcessError as e:
        logger.error("Failed to start Vaultwarden: %s", e.stderr)
        return False

    # Wait for health check
    print("Waiting for Vaultwarden to start...", end="", flush=True)
    for _ in range(30):
        try:
            result = subprocess.run(
                ["curl", "-sf", "http://127.0.0.1:8200/alive"],
                capture_output=True,
                timeout=5,
            )
            if result.returncode == 0:
                print(" ready!")
                return True
        except subprocess.TimeoutExpired:
            pass
        print(".", end="", flush=True)
        time.sleep(1)

    print(" timeout!")
    return False


def step_import_existing_secrets(bw_password: str) -> int:
    """Import existing plaintext secrets from ~/.hermes/secrets/ into Vaultwarden."""
    if not SECRETS_DIR.exists():
        return 0

    key_files = list(SECRETS_DIR.glob("*.key"))
    if not key_files:
        return 0

    print(f"\nImporting {len(key_files)} existing secrets into vault...")

    # TODO: Use Bitwarden CLI to create vault items
    # For now, document what would be imported
    imported = 0
    for key_file in key_files:
        provider = key_file.stem
        print(f"  Found: {provider} API key")
        imported += 1

    # The actual import would use:
    # bw login --apikey (or session)
    # bw create item (JSON payload with name, password, etc.)
    print(f"  {imported} secrets ready for import (Bitwarden CLI needed)")
    return imported


def run_full_setup():
    """Run the complete Full Security setup."""
    print("\n" + "=" * 50)
    print("  HERMES FULL SECURITY SETUP")
    print("=" * 50)
    print()
    print("This will set up encrypted vault storage for your agent's secrets.")
    print("You'll need pen and paper for the recovery words.")
    print()
    input("Press Enter when ready...")

    # Step 1: Mnemonic
    print("\n--- Step 1: Recovery Words ---\n")
    word_count = 12
    choice = input("Word count — 12 (standard) or 24 (paranoid)? [12]: ").strip()
    if choice == "24":
        word_count = 24

    mnemonic = step_generate_mnemonic(word_count)

    input("\nPress Enter after writing down all words...")

    # Verify first and last word
    words = mnemonic.split()
    verify = input(f"Verify — what is word #{1}? ").strip().lower()
    if verify != words[0]:
        print("Incorrect. Please start over.")
        return False

    verify = input(f"Verify — what is word #{len(words)}? ").strip().lower()
    if verify != words[-1]:
        print("Incorrect. Please start over.")
        return False

    print("Verified!")

    # Step 2: Key file
    print("\n--- Step 2: Key File (Optional) ---\n")
    keyfile_path = step_get_keyfile()

    # Step 3: Google Drive folder ID
    print("\n--- Step 3: Google Drive ---\n")
    folder_id = step_get_folder_id()

    # Step 4: Derive keys
    print("\n--- Step 4: Deriving Keys ---\n")
    kd = RookKeyDerivation(
        mnemonic=mnemonic,
        folder_id=folder_id,
        keyfile_path=keyfile_path,
    )
    print("Three independent keys derived from your recovery words.")

    # Step 5: Create LUKS container
    print("\n--- Step 5: Creating Encrypted Vault ---\n")
    container = step_create_vault(kd.luks_passphrase)

    # Step 6: Open vault
    container.open(kd.luks_passphrase)

    # Step 7: Start Vaultwarden
    print("\n--- Step 6: Starting Secret Manager ---\n")
    # Generate a random admin token
    import secrets as sec

    admin_token = sec.token_urlsafe(32)
    vw_started = step_start_vaultwarden(admin_token)

    if not vw_started:
        print("Warning: Vaultwarden failed to start. You can try again later.")

    # Step 8: Import existing secrets
    print("\n--- Step 7: Importing Existing Secrets ---\n")
    imported = step_import_existing_secrets(kd.bitwarden_password)

    # Step 9: Save state (non-sensitive)
    save_vault_state(
        {
            "setup_complete": True,
            "security_level": "full",
            "keyfile_used": keyfile_path is not None,
            "folder_id": folder_id,
            "container_path": str(container.container_path),
            "mount_point": str(container.mount_point),
            "vaultwarden_port": 8200,
            "secrets_imported": imported,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
        }
    )

    # Step 10: Wipe sensitive material from memory
    kd.wipe()
    del mnemonic

    # Done
    print("\n" + "=" * 50)
    print("  VAULT SETUP COMPLETE")
    print("=" * 50)
    print()
    print("  Your secrets are now protected by:")
    print("  - AES-256 LUKS2 encryption")
    print("  - Vaultwarden (self-hosted Bitwarden)")
    print("  - Recovery words (written down)")
    if keyfile_path:
        print(f"  - Key file: {keyfile_path}")
    print()
    print("  Seal your recovery words in an envelope.")
    print("  Store it somewhere safe — NOT next to the device.")
    print()
    return True


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Rook Vault Setup")
    sub = parser.add_subparsers(dest="command")
    sub.add_parser("setup", help="Run full security setup")
    sub.add_parser("status", help="Show vault status")
    sub.add_parser("open", help="Open vault (prompts for mnemonic)")
    sub.add_parser("close", help="Close vault")

    args = parser.parse_args()

    if args.command == "setup":
        run_full_setup()
    elif args.command == "status":
        state = load_vault_state()
        container = LuksContainer()
        print(json.dumps({**state, **container.status()}, indent=2))
    elif args.command == "open":
        # TODO: fTPM auto-unlock or PIN prompt
        print("Manual vault open not yet implemented.")
        print("Use: python3 luks_manager.py open")
    elif args.command == "close":
        container = LuksContainer()
        container.close()
        subprocess.run(
            ["docker", "compose", "-f", str(COMPOSE_FILE), "down"],
            capture_output=True,
        )
        print("Vault closed.")
    else:
        parser.print_help()