"""
LockStock Auditor — Decrypt agent conversation transcripts from the chain.

Operator tool: given an Account Key (32 bytes, never stored on server),
reconstructs full plaintext from envelope-encrypted blobs stored on the
LockStock server.

The server never sees plaintext. The auditor derives:
    Account Key → KEK (per-agent) → unwraps DEK → blob_key (per-blob) → decrypt

Usage:
    lockstock-audit --account-key <hex> --agent-id <id> --api-key <key> [options]

    # Decrypt last 20 entries
    lockstock-audit --account-key 85d7a5... --agent-id agent_abc --api-key lsk_admin_...

    # JSON output for machine processing
    lockstock-audit --account-key 85d7a5... --agent-id agent_abc --api-key lsk_admin_... --format json

    # Specific sequence range
    lockstock-audit ... --from-seq 17 --to-seq 18
"""

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

import httpx

from .keyring import derive_kek, unwrap_dek
from .encrypt import decrypt_payload_envelope

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


def fetch_audit_entries(
    server_url: str,
    agent_id: str,
    api_key: str,
    limit: int = DEFAULT_LIMIT,
) -> List[Dict[str, Any]]:
    """Fetch audit log entries from LockStock server."""
    headers = {"Authorization": f"Bearer {api_key}"}
    r = httpx.get(
        f"{server_url}/api/guard/audit",
        params={"agent_id": agent_id, "limit": limit},
        headers=headers,
        timeout=30.0,
    )
    r.raise_for_status()
    return r.json()


def fetch_wrapped_dek(
    server_url: str,
    agent_id: str,
    api_key: str,
    key_version: Optional[int] = None,
) -> Dict[str, Any]:
    """Fetch wrapped DEK from server."""
    headers = {"Authorization": f"Bearer {api_key}"}
    version_path = str(key_version) if key_version is not None else "latest"
    r = httpx.get(
        f"{server_url}/v1/agents/{agent_id}/keys/{version_path}",
        headers=headers,
        timeout=30.0,
    )
    r.raise_for_status()
    return r.json()


def fetch_payload_blob(
    server_url: str,
    payload_hash: str,
    api_key: str,
) -> Optional[Dict[str, Any]]:
    """Fetch encrypted payload blob by content address."""
    headers = {"Authorization": f"Bearer {api_key}"}
    r = httpx.get(
        f"{server_url}/v1/payload/{payload_hash}",
        headers=headers,
        timeout=30.0,
    )
    if r.status_code == 404:
        return None
    r.raise_for_status()
    return r.json()


def decrypt_blob(
    blob: Dict[str, Any],
    dek: bytes,
) -> bytes:
    """Decrypt an encrypted payload blob using the DEK."""
    ciphertext = base64.b64decode(blob["ciphertext_b64"])
    stamp_hash = blob["stamp_hash"]
    payload_hash = blob["payload_hash"]
    return decrypt_payload_envelope(ciphertext, stamp_hash, payload_hash, dek)


def audit_agent(
    account_key: bytes,
    agent_id: str,
    api_key: str,
    server_url: str = DEFAULT_SERVER,
    limit: int = DEFAULT_LIMIT,
    from_seq: Optional[int] = None,
    to_seq: Optional[int] = None,
) -> List[Dict[str, Any]]:
    """
    Audit an agent's conversation history by decrypting envelope-encrypted blobs.

    Returns a list of transcript entries, each containing:
        - sequence: chain sequence number
        - task: action type (CHATCOMPLETION, CHATCOMPLETIONRESPONSE, etc.)
        - timestamp: server timestamp
        - payload_hash: content address
        - stamp_hash: chain position hash
        - key_version: 0 = chain-as-key, >=1 = envelope
        - plaintext_length: decrypted size in bytes
        - content: decoded plaintext (bencoded dict as string, or raw bytes repr)
        - decrypted: True if successfully decrypted, False otherwise
        - error: error message if decryption failed
    """
    # Step 1: Fetch audit entries
    entries = fetch_audit_entries(server_url, agent_id, api_key, limit)

    # Step 2: Derive KEK and unwrap DEK
    kek = derive_kek(account_key, agent_id)
    dek_data = fetch_wrapped_dek(server_url, agent_id, api_key)
    wrapped_dek = base64.b64decode(dek_data["wrapped_dek_b64"])
    dek = unwrap_dek(kek, wrapped_dek)
    dek_version = dek_data.get("key_version", 1)

    # Step 3: Walk entries, fetch and decrypt blobs
    transcript = []
    for entry in reversed(entries):  # oldest first
        details = entry.get("details", {})
        seq = details.get("sequence_id")
        payload_hash = details.get("payload_hash")
        task = entry.get("action", "")
        timestamp = entry.get("timestamp", "")

        # Filter by sequence range (skip entries without sequence when filtering)
        if from_seq is not None and (seq is None or seq < from_seq):
            continue
        if to_seq is not None and (seq is None or seq > to_seq):
            continue

        record = {
            "sequence": seq,
            "task": task,
            "timestamp": timestamp,
            "payload_hash": payload_hash,
            "decrypted": False,
        }

        if not payload_hash:
            record["error"] = "no payload_hash in chain entry"
            transcript.append(record)
            continue

        # Fetch blob
        blob = fetch_payload_blob(server_url, payload_hash, api_key)
        if blob is None:
            record["error"] = "blob not found on server"
            transcript.append(record)
            continue

        record["stamp_hash"] = blob.get("stamp_hash")
        record["key_version"] = blob.get("key_version", 0)

        # Decrypt
        try:
            plaintext = decrypt_blob(blob, dek)
            record["plaintext_length"] = len(plaintext)
            record["content"] = plaintext.decode("utf-8", errors="replace")
            record["decrypted"] = True
        except Exception as e:
            record["error"] = str(e)

        transcript.append(record)

    return transcript


def format_text(transcript: List[Dict[str, Any]], agent_id: str) -> str:
    """Format transcript as human-readable text."""
    lines = []
    lines.append(f"LockStock Audit Transcript — {agent_id}")
    lines.append("=" * 60)
    lines.append("")

    for entry in transcript:
        seq = entry.get("sequence", "?")
        task = entry.get("task", "UNKNOWN")
        ts = entry.get("timestamp", "")
        decrypted = entry.get("decrypted", False)

        lines.append(f"[seq {seq}] {task}  ({ts})")

        if decrypted:
            content = entry.get("content", "")
            kv = entry.get("key_version", 0)
            plen = entry.get("plaintext_length", 0)
            lines.append(f"  key_version={kv}  plaintext={plen} bytes")
            lines.append("")

            # Try to parse bencoded content for prettier output
            try:
                parsed = _extract_readable(content)
                if parsed:
                    lines.append(f"  {parsed}")
                else:
                    lines.append(f"  {content[:500]}")
            except Exception:
                lines.append(f"  {content[:500]}")
        else:
            error = entry.get("error", "unknown error")
            lines.append(f"  [NOT DECRYPTED: {error}]")

        lines.append("")
        lines.append("-" * 60)
        lines.append("")

    return "\n".join(lines)


def _extract_readable(bencoded_content: str) -> Optional[str]:
    """
    Extract human-readable text from bencoded content.

    Bencoded dicts have patterns like:
        7:content25:Trinity, what do you see?
        8:messages...7:content25:Trinity, what do you see?

    This extracts the longest 'content' value as the readable message.
    """
    # Find all content values in bencoded string
    # Pattern: N:contentM:actual_text where N=7 and M is the length
    import re
    matches = []
    # Look for content field values: digit(s):content followed by digit(s):text
    for m in re.finditer(r'7:content(\d+):', bencoded_content):
        length = int(m.group(1))
        start = m.end()
        text = bencoded_content[start:start + length]
        if text and len(text) > 5:  # skip trivially short
            matches.append(text)

    if matches:
        # Return the longest content (usually the actual message)
        return max(matches, key=len)
    return None


def main(argv: Optional[List[str]] = None) -> int:
    """CLI entry point for lockstock-audit."""
    parser = argparse.ArgumentParser(
        prog="lockstock-audit",
        description="Decrypt agent conversation transcripts from the LockStock chain.",
        epilog="The server never sees plaintext. Only the Account Key holder can audit.",
    )
    parser.add_argument(
        "--account-key", required=True,
        help="Account Key as hex string (64 chars = 32 bytes)",
    )
    parser.add_argument(
        "--agent-id", required=True,
        help="Agent ID to audit",
    )
    parser.add_argument(
        "--api-key", required=True,
        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(
        "--limit", type=int, default=DEFAULT_LIMIT,
        help=f"Maximum audit entries to fetch (default: {DEFAULT_LIMIT})",
    )
    parser.add_argument(
        "--from-seq", type=int, default=None,
        help="Start from this sequence number (inclusive)",
    )
    parser.add_argument(
        "--to-seq", type=int, default=None,
        help="End at this sequence number (inclusive)",
    )
    parser.add_argument(
        "--format", choices=["text", "json"], default="text",
        dest="output_format",
        help="Output format (default: text)",
    )

    args = parser.parse_args(argv)

    # Validate account key
    try:
        account_key = bytes.fromhex(args.account_key)
        if len(account_key) != 32:
            print(f"Error: Account Key must be 32 bytes (64 hex chars), got {len(account_key)}", file=sys.stderr)
            return 1
    except ValueError:
        print("Error: Account Key must be valid hex", file=sys.stderr)
        return 1

    # Run audit
    try:
        transcript = audit_agent(
            account_key=account_key,
            agent_id=args.agent_id,
            api_key=args.api_key,
            server_url=args.server_url,
            limit=args.limit,
            from_seq=args.from_seq,
            to_seq=args.to_seq,
        )
    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(transcript, indent=2))
    else:
        print(format_text(transcript, args.agent_id))

    # Summary
    total = len(transcript)
    decrypted = sum(1 for e in transcript if e.get("decrypted"))
    missing = sum(1 for e in transcript if e.get("error") == "blob not found on server")
    failed = total - decrypted - missing
    if args.output_format == "text":
        print(f"Summary: {decrypted}/{total} entries decrypted"
              f"{f', {missing} blobs missing' if missing else ''}"
              f"{f', {failed} failed' if failed else ''}")

    return 0


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