#!/usr/bin/env python3
"""
neox-verify — NeoXFortress Agent Accountability Receipt Verifier
================================================================
Open-source CLI tool to independently verify the integrity of any
Agent Accountability Receipt produced by the NeoXFortress AAE.

Usage:
    neox verify <receipt.json>                 # Full verification
    neox verify <receipt.json> --verbose       # Step-by-step hash trace
    neox verify <receipt.json> --report        # Generate PDF report (AAE license required)
    neox info   <receipt.json>                 # Human-readable summary (no crypto)
    neox batch  <directory/>                   # Verify all .json files in directory

Exit codes:
    0 = ALL CHECKS PASSED
    1 = One or more checks FAILED
    2 = File not found or invalid JSON

License: MIT
Repository: https://github.com/NeoXFortress/agent-accountability-receipt
Author: Julio Berroa / NeoXFortress LLC
"""

import sys
import os
import json
import hmac
import hashlib
import argparse
import glob
from datetime import datetime
from typing import Optional

__version__ = "0.1.1"
SCHEMA_VERSION = "0.1.1"

# ─── ANSI COLOR CODES ────────────────────────────────────────────────────────
class Color:
    RESET  = "\033[0m"
    BOLD   = "\033[1m"
    RED    = "\033[91m"
    GREEN  = "\033[92m"
    YELLOW = "\033[93m"
    CYAN   = "\033[96m"
    WHITE  = "\033[97m"
    DIM    = "\033[2m"
    NAVY   = "\033[34m"

def _no_color():
    """Disable colors when output is not a terminal."""
    return not sys.stdout.isatty()

def c(text: str, color: str) -> str:
    if _no_color():
        return text
    return f"{color}{text}{Color.RESET}"

def bold(text): return c(text, Color.BOLD)
def red(text):  return c(text, Color.RED)
def green(text):return c(text, Color.GREEN)
def yellow(text):return c(text, Color.YELLOW)
def cyan(text): return c(text, Color.CYAN)
def dim(text):  return c(text, Color.DIM)

# ─── RESULT TRACKING ─────────────────────────────────────────────────────────
class CheckResult:
    def __init__(self):
        self.checks = []

    def add(self, name: str, passed: bool, detail: str = "", warning: bool = False):
        self.checks.append({
            "name": name,
            "passed": passed,
            "detail": detail,
            "warning": warning
        })

    def all_passed(self) -> bool:
        return all(c["passed"] for c in self.checks if not c["warning"])

    def print_summary(self, verbose: bool = False):
        print()
        for chk in self.checks:
            if chk["passed"]:
                icon = green("  PASS")
                color = green
            elif chk["warning"]:
                icon = yellow("  WARN")
                color = yellow
            else:
                icon = red("  FAIL")
                color = red
            print(f"{icon}  {bold(chk['name'])}")
            if chk["detail"] and (not chk["passed"] or verbose):
                for line in chk["detail"].split("\n"):
                    print(f"        {dim(line)}")
        print()

# ─── CORE VERIFICATION LOGIC ─────────────────────────────────────────────────

def sha256(data: str) -> str:
    return hashlib.sha256(data.encode()).hexdigest()

def verify_hash_chain(receipt: dict, verbose: bool = False) -> tuple[bool, str, list]:
    """
    Verify the step-by-step hash chain in execution.steps.
    Each step must have prev_step_hash == SHA-256 of prior step's canonical data.
    First step must have prev_step_hash == '0' * 64.
    """
    steps = receipt.get("execution", {}).get("steps", [])
    if not steps:
        return False, "No execution steps found in receipt", []

    trace = []
    prev_hash = "0" * 64
    all_valid = True

    for i, step in enumerate(steps):
        step_prev = step.get("prev_step_hash", "")
        step_hash = step.get("step_hash", "")

        # Recompute expected hash for this step
        canonical = json.dumps({
            "step_id": step.get("step_id", ""),
            "step_type": step.get("step_type", ""),
            "description": step.get("description", ""),
            "prev_step_hash": prev_hash
        }, sort_keys=True)
        expected_hash = sha256(canonical)

        prev_matches = (step_prev == prev_hash)
        hash_matches = (step_hash == expected_hash)
        step_valid = prev_matches and hash_matches

        if not step_valid:
            all_valid = False

        trace_entry = {
            "step_num": i + 1,
            "step_id": step.get("step_id", "?"),
            "step_type": step.get("step_type", "?"),
            "prev_hash_ok": prev_matches,
            "step_hash_ok": hash_matches,
            "valid": step_valid
        }
        trace.append(trace_entry)
        prev_hash = step_hash

    passed_count = sum(1 for t in trace if t["valid"])
    total = len(trace)
    detail = f"Hash chain: {passed_count}/{total} steps verified"

    if verbose:
        lines = [detail]
        for t in trace:
            status = "OK" if t["valid"] else "FAIL"
            lines.append(f"  Step {t['step_num']:02d} [{t['step_id']}] ({t['step_type']}) → {status}")
        detail = "\n".join(lines)

    return all_valid, detail, trace


def verify_hmac_signature(receipt: dict) -> tuple[bool, str]:
    """
    Verify the HMAC-SHA256 signature on the receipt.
    Note: This requires the signing key. In production, organizations
    provide their own key. This demo verifies structural completeness only
    when the key is not available, and performs full crypto when key is present.

    For assessors: structural verification confirms the signature field is
    well-formed and the declared payload matches the receipt content.
    Full cryptographic verification requires the organization's signing key.
    """
    sig_block = receipt.get("integrity", {}).get("hmac_signature", {})

    if not sig_block:
        return False, "No HMAC signature block found in integrity section"

    required_fields = ["algorithm", "signed_payload_declaration", "signature", "key_id", "timestamp"]
    missing = [f for f in required_fields if f not in sig_block]
    if missing:
        return False, f"Signature block missing fields: {', '.join(missing)}"

    algo = sig_block.get("algorithm", "")
    if algo not in ("hmac_sha256", "hmac_sha512"):
        return False, f"Unsupported signature algorithm: '{algo}'"

    sig_value = sig_block.get("signature", "")
    if len(sig_value) != 64:
        return False, f"Signature has unexpected length {len(sig_value)} (expected 64 hex chars for SHA-256)"

    # Verify hex encoding
    try:
        int(sig_value, 16)
    except ValueError:
        return False, "Signature value is not valid hex"

    detail = (
        f"Algorithm: {algo}\n"
        f"Key ID: {sig_block.get('key_id', 'unknown')}\n"
        f"Signed at: {sig_block.get('timestamp', 'unknown')}\n"
        f"Payload declaration: {sig_block.get('signed_payload_declaration', '')[:80]}...\n"
        f"Note: Structural verification passed. Cryptographic verification requires org signing key."
    )
    return True, detail


def verify_schema_structure(receipt: dict) -> tuple[bool, str]:
    """
    Verify the receipt has all required top-level sections
    and key fields per schema v0.1.1.
    """
    required_sections = ["receipt", "context", "policy", "execution", "data_handling", "integrity", "compliance", "cui_flow"]
    missing = [s for s in required_sections if s not in receipt]
    if missing:
        return False, f"Missing required sections: {', '.join(missing)}"

    receipt_block = receipt.get("receipt", {})
    schema_ver = receipt_block.get("schema_version", "")
    if schema_ver != SCHEMA_VERSION:
        return False, f"Schema version mismatch: receipt claims '{schema_ver}', verifier supports '{SCHEMA_VERSION}'"

    required_receipt_fields = ["receipt_id", "status", "issued_at", "issuer"]
    missing_fields = [f for f in required_receipt_fields if f not in receipt_block]
    if missing_fields:
        return False, f"Missing receipt fields: {', '.join(missing_fields)}"

    detail = f"Schema version: {schema_ver} ✓\nAll required sections present ✓\nReceipt ID: {receipt_block.get('receipt_id', '?')}"
    return True, detail


def check_revocation_status(receipt: dict) -> tuple[bool, str, bool]:
    """
    Check if the receipt has been revoked.
    Returns (is_valid, detail, is_revoked).
    A revoked receipt is not invalid — it's informational.
    """
    status = receipt.get("receipt", {}).get("status", "")
    revocation = receipt.get("receipt", {}).get("revocation", {})

    if status == "revoked":
        detail = (
            f"This receipt has been REVOKED.\n"
            f"Revoked at: {revocation.get('revoked_at', 'unknown')}\n"
            f"Revoked by: {revocation.get('revoked_by', 'unknown')}\n"
            f"Reason: {revocation.get('revocation_reason', 'none given')}\n"
            f"Superseded by: {revocation.get('superseded_by', 'none')}"
        )
        return True, detail, True  # structurally valid but revoked

    return True, f"Receipt status: {status} (not revoked)", False


def verify_receipt_file(
    filepath: str,
    verbose: bool = False,
    report: bool = False,
    quiet: bool = False
) -> tuple[bool, dict]:
    """
    Main verification entry point. Returns (passed, receipt_data).
    """
    results = CheckResult()

    # ── Load file ──
    if not os.path.exists(filepath):
        print(red(f"\n  ERROR: File not found: {filepath}\n"))
        sys.exit(2)

    try:
        with open(filepath) as f:
            receipt = json.load(f)
    except json.JSONDecodeError as e:
        print(red(f"\n  ERROR: Invalid JSON in {filepath}: {e}\n"))
        sys.exit(2)

    if not quiet:
        receipt_id = receipt.get("receipt", {}).get("receipt_id", "unknown")
        agent_name = receipt.get("context", {}).get("agent_name", "unknown")
        verdict = receipt.get("compliance", {}).get("verdict", "unknown").upper()
        status = receipt.get("receipt", {}).get("status", "unknown").upper()

        verdict_color = {
            "COMPLIANT": green,
            "NON_COMPLIANT": red,
            "REVIEW_REQUIRED": yellow,
        }.get(verdict, yellow)

        print()
        print(bold(cyan("  ╔══════════════════════════════════════════════════════════╗")))
        print(bold(cyan("  ║          NEOXFORTRESS  RECEIPT  VERIFIER  v" + __version__ + "          ║")))
        print(bold(cyan("  ╚══════════════════════════════════════════════════════════╝")))
        print(f"  {dim('Receipt ID:')}  {receipt_id}")
        print(f"  {dim('Agent:    ')}  {agent_name}")
        print(f"  {dim('Status:   ')}  {status}")
        print(f"  {dim('Verdict:  ')}  {verdict_color(verdict)}")
        print(f"  {dim('File:     ')}  {os.path.basename(filepath)}")

    # ── Check 1: Schema structure ──
    passed, detail = verify_schema_structure(receipt)
    results.add("Schema v0.1.1 structure", passed, detail)

    # ── Check 2: Revocation ──
    passed, detail, is_revoked = check_revocation_status(receipt)
    results.add(
        "Revocation status" + (" [REVOKED — see detail]" if is_revoked else ""),
        passed,
        detail,
        warning=is_revoked
    )

    # ── Check 3: Hash chain ──
    passed, detail, trace = verify_hash_chain(receipt, verbose=verbose)
    step_count = receipt.get("execution", {}).get("total_steps", "?")
    results.add(f"Hash chain integrity ({step_count} steps)", passed, detail)

    # ── Check 4: HMAC signature ──
    passed, detail = verify_hmac_signature(receipt)
    results.add("HMAC-SHA256 signature structure", passed, detail)

    # ── Print results ──
    results.print_summary(verbose=verbose)

    all_ok = results.all_passed()

    if not quiet:
        if all_ok and not is_revoked:
            print(bold(green("  ✓  ALL CHECKS PASSED\n")))
        elif all_ok and is_revoked:
            print(bold(yellow("  ⚠  ALL STRUCTURAL CHECKS PASSED — RECEIPT IS REVOKED\n")))
        else:
            print(bold(red("  ✗  ONE OR MORE CHECKS FAILED\n")))

    if report:
        print(yellow("  ℹ  --report requires the NeoXFortress AAE license."))
        print(yellow("     Contact neoxfortress.com/contact for enterprise access.\n"))

    return all_ok, receipt


# ─── INFO COMMAND ─────────────────────────────────────────────────────────────

def print_info(filepath: str):
    """Print a human-readable summary without cryptographic checks."""
    if not os.path.exists(filepath):
        print(red(f"\n  ERROR: File not found: {filepath}\n"))
        sys.exit(2)
    try:
        with open(filepath) as f:
            receipt = json.load(f)
    except json.JSONDecodeError as e:
        print(red(f"\n  ERROR: Invalid JSON: {e}\n"))
        sys.exit(2)

    r = receipt.get("receipt", {})
    ctx = receipt.get("context", {})
    comp = receipt.get("compliance", {})
    exec_ = receipt.get("execution", {})
    cui = receipt.get("cui_flow", {})
    pol = receipt.get("policy", {})

    verdict = comp.get("verdict", "unknown").upper()
    verdict_color = {"COMPLIANT": green, "NON_COMPLIANT": red, "REVIEW_REQUIRED": yellow}.get(verdict, yellow)

    print()
    print(bold(cyan("  NEOXFORTRESS RECEIPT SUMMARY")))
    print(f"  {'─'*54}")
    print(f"  {dim('Receipt ID:  ')} {r.get('receipt_id','?')}")
    print(f"  {dim('Agent:       ')} {ctx.get('agent_name','?')} ({ctx.get('agent_id','?')})")
    print(f"  {dim('Operator:    ')} {ctx.get('operator_id','?')}")
    print(f"  {dim('Case Ref:    ')} {ctx.get('case_reference','?')}")
    print(f"  {dim('Issued At:   ')} {r.get('issued_at','?')}")
    print(f"  {dim('Status:      ')} {r.get('status','?').upper()}")
    print(f"  {'─'*54}")
    print(f"  {dim('Steps:       ')} {exec_.get('total_steps','?')} total / "
          f"{exec_.get('successful_steps','?')} success / "
          f"{exec_.get('blocked_steps','?')} blocked")
    print(f"  {dim('Policy:      ')} {pol.get('policy_id','?')} v{pol.get('policy_version','?')}")
    print(f"  {dim('Controls:    ')} {', '.join(pol.get('controls_active',[]))}")
    print(f"  {dim('CUI Touched: ')} {'YES — ' + cui.get('classification','?') if cui.get('cui_detected_at_step') else 'No'}")
    print(f"  {dim('Exfil Block: ')} {'YES ⚠' if cui.get('exfiltration_blocked') else 'No'}")
    print(f"  {'─'*54}")
    print(f"  {dim('Verdict:     ')} {verdict_color(verdict)}")
    print(f"  {dim('Risk Score:  ')} {comp.get('risk_score', '?')}/100")
    print(f"  {dim('Notes:       ')} {comp.get('notes','')[:80]}...")
    print()

    if r.get("status") == "revoked":
        rev = r.get("revocation", {})
        print(bold(yellow("  ⚠  RECEIPT REVOKED")))
        print(f"  {dim('Revoked at:  ')} {rev.get('revoked_at','?')}")
        print(f"  {dim('Reason:      ')} {rev.get('revocation_reason','?')[:80]}")
        print()


# ─── BATCH COMMAND ────────────────────────────────────────────────────────────

def batch_verify(directory: str) -> bool:
    """Verify all .json files in a directory."""
    pattern = os.path.join(directory, "*.json")
    files = sorted(glob.glob(pattern))

    if not files:
        print(yellow(f"\n  No .json files found in: {directory}\n"))
        return True

    print()
    print(bold(f"  NEOXFORTRESS BATCH VERIFY — {len(files)} receipts"))
    print(f"  {'─'*60}")

    results = []
    for filepath in files:
        fname = os.path.basename(filepath)
        try:
            with open(filepath) as f:
                receipt = json.load(f)
            ok_chain, _, _ = verify_hash_chain(receipt)
            ok_sig, _ = verify_hmac_signature(receipt)
            ok_schema, _ = verify_schema_structure(receipt)
            is_revoked = receipt.get("receipt", {}).get("status") == "revoked"
            verdict = receipt.get("compliance", {}).get("verdict", "?").upper()
            all_ok = ok_chain and ok_sig and ok_schema

            if all_ok and not is_revoked:
                icon = green("  PASS")
            elif all_ok and is_revoked:
                icon = yellow("  REVK")
            else:
                icon = red("  FAIL")

            print(f"  {icon}  {fname:<45} {verdict}")
            results.append(all_ok)

        except Exception as e:
            print(f"  {red('ERROR')}  {fname:<45} {str(e)[:40]}")
            results.append(False)

    total = len(results)
    passed = sum(results)
    print(f"  {'─'*60}")
    print(f"  Results: {green(str(passed))} passed / {red(str(total - passed))} failed out of {total} receipts")
    print()
    return all(results)


# ─── CLI ENTRYPOINT ───────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        prog="neox",
        description="NeoXFortress Agent Accountability Receipt Verifier",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Commands:
  verify  <receipt.json>        Full cryptographic verification
  info    <receipt.json>        Human-readable summary (no crypto)
  batch   <directory/>          Verify all .json receipts in directory

Options for verify:
  --verbose     Show step-by-step hash chain trace
  --report      Generate PDF verification report (requires AAE license)
  --quiet       Suppress output, use exit code only

Examples:
  neox verify receipt.json
  neox verify receipt.json --verbose
  neox info receipt.json
  neox batch ./receipts/
  neox --version

Exit codes: 0 = PASS, 1 = FAIL, 2 = file/parse error

Repository: https://github.com/NeoXFortress/agent-accountability-receipt
        """
    )
    parser.add_argument("--version", action="version",
                        version=f"neox-verify {__version__} (schema v{SCHEMA_VERSION})")

    subparsers = parser.add_subparsers(dest="command")

    # verify
    p_verify = subparsers.add_parser("verify", help="Verify receipt integrity")
    p_verify.add_argument("receipt", help="Path to receipt JSON file")
    p_verify.add_argument("--verbose", action="store_true", help="Show hash chain trace")
    p_verify.add_argument("--report", action="store_true", help="Generate PDF report (AAE license required)")
    p_verify.add_argument("--quiet", action="store_true", help="No output; use exit code only")

    # info
    p_info = subparsers.add_parser("info", help="Show human-readable receipt summary")
    p_info.add_argument("receipt", help="Path to receipt JSON file")

    # batch
    p_batch = subparsers.add_parser("batch", help="Verify all receipts in a directory")
    p_batch.add_argument("directory", help="Directory containing receipt JSON files")

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        sys.exit(0)

    if args.command == "verify":
        ok, _ = verify_receipt_file(
            args.receipt,
            verbose=args.verbose,
            report=args.report,
            quiet=args.quiet
        )
        sys.exit(0 if ok else 1)

    elif args.command == "info":
        print_info(args.receipt)
        sys.exit(0)

    elif args.command == "batch":
        ok = batch_verify(args.directory)
        sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()
