"""
LockStock Account Key Rotation — Re-wrap all agent DEKs under a new Account Key.

The operator rotates their Account Key by:
1. Fetching all wrapped DEKs via GET /v1/operators/keys
2. For each agent: unwrap DEK with old KEK, re-wrap with new KEK
3. PUT each re-wrapped DEK back to the server

The server never sees the plaintext DEK. It stores opaque wrapped bytes
that changed shape but still protect the same underlying key.

Usage:
    lockstock-rotate-key --old-key <hex> --new-key <hex> --api-key <key> [options]

    # Rotate all agents to new Account Key
    lockstock-rotate-key --old-key 85d7a5... --new-key aa11bb22... --api-key lsk_admin_...

    # Dry run (show what would change without writing)
    lockstock-rotate-key --old-key 85d7a5... --new-key aa11bb22... --api-key lsk_admin_... --dry-run

    # Generate a new Account Key
    lockstock-rotate-key --generate-key
"""

import argparse
import base64
import json
import sys
from typing import Any, Dict, List, Optional

import httpx

from .keyring import (
    derive_kek,
    generate_account_key,
    unwrap_dek,
    wrap_dek,
)

DEFAULT_SERVER = "https://lockstock-api-i9kp.onrender.com"


def fetch_operator_keys(
    server_url: str,
    api_key: str,
) -> Dict[str, Any]:
    """Fetch all wrapped DEKs for the authenticated operator."""
    headers = {"Authorization": f"Bearer {api_key}"}
    r = httpx.get(
        f"{server_url}/v1/operators/keys",
        headers=headers,
        timeout=30.0,
    )
    r.raise_for_status()
    return r.json()


def store_wrapped_dek(
    server_url: str,
    agent_id: str,
    key_version: int,
    wrapped_dek_b64: str,
    api_key: str,
) -> Dict[str, Any]:
    """Store a re-wrapped DEK for an agent."""
    headers = {"Authorization": f"Bearer {api_key}"}
    r = httpx.put(
        f"{server_url}/v1/agents/{agent_id}/keys/{key_version}",
        json={"wrapped_dek_b64": wrapped_dek_b64},
        headers=headers,
        timeout=30.0,
    )
    r.raise_for_status()
    return r.json()


def rotate_account_key(
    old_account_key: bytes,
    new_account_key: bytes,
    api_key: str,
    server_url: str = DEFAULT_SERVER,
    dry_run: bool = False,
) -> List[Dict[str, Any]]:
    """
    Rotate all agent DEKs from old Account Key to new Account Key.

    For each agent owned by the operator:
    1. Derive old KEK from (old_account_key, agent_id)
    2. Unwrap DEK using old KEK
    3. Derive new KEK from (new_account_key, agent_id)
    4. Re-wrap DEK using new KEK
    5. PUT new wrapped DEK to server (same key_version, new wrapping)

    Returns a list of results per agent, each containing:
        - agent_id: the agent's identifier
        - key_version: the DEK version that was re-wrapped
        - status: "rotated", "dry_run", or "failed"
        - error: error message if failed
    """
    # Step 1: Fetch all wrapped DEKs
    data = fetch_operator_keys(server_url, api_key)
    agents = data.get("agents", [])
    operator_id = data.get("operator_id", "unknown")

    results = []

    for agent_entry in agents:
        agent_id = agent_entry["agent_id"]
        key_version = agent_entry["key_version"]
        wrapped_b64 = agent_entry["wrapped_dek_b64"]

        result = {
            "agent_id": agent_id,
            "key_version": key_version,
        }

        try:
            # Step 2: Unwrap with old key
            old_kek = derive_kek(old_account_key, agent_id)
            wrapped_bytes = base64.b64decode(wrapped_b64)
            dek = unwrap_dek(old_kek, wrapped_bytes)

            # Step 3: Re-wrap with new key
            new_kek = derive_kek(new_account_key, agent_id)
            new_wrapped = wrap_dek(new_kek, dek)
            new_wrapped_b64 = base64.b64encode(new_wrapped).decode("ascii")

            if dry_run:
                result["status"] = "dry_run"
                result["old_wrapped_prefix"] = wrapped_b64[:16] + "..."
                result["new_wrapped_prefix"] = new_wrapped_b64[:16] + "..."
            else:
                # Step 4: Store re-wrapped DEK
                store_wrapped_dek(
                    server_url, agent_id, key_version, new_wrapped_b64, api_key,
                )
                result["status"] = "rotated"

        except Exception as e:
            result["status"] = "failed"
            result["error"] = str(e)

        results.append(result)

    return results


def main(argv: Optional[List[str]] = None) -> int:
    """CLI entry point for lockstock-rotate-key."""
    parser = argparse.ArgumentParser(
        prog="lockstock-rotate-key",
        description="Rotate Account Key — re-wrap all agent DEKs under a new key.",
        epilog="The server never sees the plaintext DEK during rotation.",
    )
    parser.add_argument(
        "--old-key",
        help="Current Account Key as hex string (64 chars = 32 bytes)",
    )
    parser.add_argument(
        "--new-key",
        help="New Account Key as hex string (64 chars = 32 bytes)",
    )
    parser.add_argument(
        "--api-key",
        help="LockStock API key (Bearer token)",
    )
    parser.add_argument(
        "--server-url", default=DEFAULT_SERVER,
        help=f"LockStock server URL (default: {DEFAULT_SERVER})",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Show what would change without writing to server",
    )
    parser.add_argument(
        "--generate-key", action="store_true",
        help="Generate a new Account Key and exit",
    )
    parser.add_argument(
        "--format", choices=["text", "json"], default="text",
        dest="output_format",
        help="Output format (default: text)",
    )

    args = parser.parse_args(argv)

    # Generate key mode
    if args.generate_key:
        key = generate_account_key()
        print(key.hex())
        return 0

    # Validate required args for rotation
    if not args.old_key or not args.new_key or not args.api_key:
        parser.error("--old-key, --new-key, and --api-key are required for rotation")
        return 1

    # Validate keys
    for name, value in [("old-key", args.old_key), ("new-key", args.new_key)]:
        try:
            key_bytes = bytes.fromhex(value)
            if len(key_bytes) != 32:
                print(f"Error: --{name} must be 32 bytes (64 hex chars), got {len(key_bytes)}", file=sys.stderr)
                return 1
        except ValueError:
            print(f"Error: --{name} must be valid hex", file=sys.stderr)
            return 1

    old_key = bytes.fromhex(args.old_key)
    new_key = bytes.fromhex(args.new_key)

    if old_key == new_key:
        print("Error: old and new Account Keys are identical", file=sys.stderr)
        return 1

    # Run rotation
    try:
        results = rotate_account_key(
            old_account_key=old_key,
            new_account_key=new_key,
            api_key=args.api_key,
            server_url=args.server_url,
            dry_run=args.dry_run,
        )
    except httpx.HTTPStatusError as e:
        print(f"HTTP error: {e.response.status_code} {e.response.text}", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        return 1

    # Output
    if args.output_format == "json":
        print(json.dumps(results, indent=2))
    else:
        mode = "DRY RUN" if args.dry_run else "ROTATION"
        print(f"Account Key {mode}")
        print("=" * 50)
        for r in results:
            status = r["status"].upper()
            print(f"  {r['agent_id']} v{r['key_version']}: {status}")
            if r.get("error"):
                print(f"    Error: {r['error']}")
        print()
        rotated = sum(1 for r in results if r["status"] in ("rotated", "dry_run"))
        failed = sum(1 for r in results if r["status"] == "failed")
        print(f"Total: {len(results)} agents, {rotated} {'would rotate' if args.dry_run else 'rotated'}"
              f"{f', {failed} failed' if failed else ''}")

    return 0


if __name__ == "__main__":
    sys.exit(main())
