demo.py

"""
demo.py
-------
Interactive demo of the Smart Contract Audit RL Environment.
Shows Task 1 end-to-end with a human-readable step-by-step walkthrough.

Usage:
  python demo.py              # interactive mode
  python demo.py --auto       # auto-run with built-in demo agent (no input needed)
  python demo.py --auto --seed 42

Great for hackathon demos — run this live to show the environment in action.
"""

import argparse
import sys
import textwrap
import time

# ─────────────────────────────────────────────────────────────────────────────
# Imports
# ─────────────────────────────────────────────────────────────────────────────
from tasks.task1.environment import Task1Environment
from env.schemas import Action, ActionType

# ─────────────────────────────────────────────────────────────────────────────
# ANSI colours (falls back gracefully on Windows)
# ─────────────────────────────────────────────────────────────────────────────
try:
    import os
    if os.name == "nt":
        raise ImportError
    BOLD   = "\033[1m"
    DIM    = "\033[2m"
    GREEN  = "\033[92m"
    YELLOW = "\033[93m"
    RED    = "\033[91m"
    CYAN   = "\033[96m"
    RESET  = "\033[0m"
except ImportError:
    BOLD = DIM = GREEN = YELLOW = RED = CYAN = RESET = ""

DIVIDER = f"{DIM}{'─' * 64}{RESET}"


# ─────────────────────────────────────────────────────────────────────────────
# Pretty printers
# ─────────────────────────────────────────────────────────────────────────────

def banner():
    print()
    print(f"{BOLD}{CYAN}╔══════════════════════════════════════════════════════════╗")
    print(f"║   Smart Contract Audit RL Environment  ·  Task 1 Demo    ║")
    print(f"╚══════════════════════════════════════════════════════════╝{RESET}")
    print()


def print_observation(obs):
    print(DIVIDER)
    print(f"{BOLD}Contract :{RESET} {obs.contract_name}")
    print(f"{BOLD}Desc     :{RESET} {textwrap.fill(obs.contract_description, 72, subsequent_indent=' ' * 11)}")
    print(f"{BOLD}Step     :{RESET} {obs.step_count}   "
          f"{BOLD}Reward   :{RESET} {obs.cumulative_reward:+.2f}")
    if obs.last_action:
        colour = GREEN if obs.cumulative_reward >= 0 else YELLOW
        result = obs.last_action_result or ""
        print(f"{BOLD}Last     :{RESET} [{obs.last_action}]")
        for line in textwrap.wrap(result, 72):
            print(f"           {colour}{line}{RESET}")
    print(DIVIDER)


def print_action_menu():
    actions = [
        ("1", "list_functions",       "{}",                                         "List all functions"),
        ("2", "get_function_code",    '{"function_name": "???"}',                   "Get function source code"),
        ("3", "get_function_summary", '{"function_name": "???"}',                   "Get NatSpec comment"),
        ("4", "get_file_metadata",    "{}",                                         "Get file metadata"),
        ("5", "get_state_variable",   '{"variable_name": "???"}',                   "Get state variable"),
        ("6", "get_call_graph",       "{}",                                         "Get call graph"),
        ("7", "submit",               '{"function_name":"???","vulnerability_type":"???"}', "Submit answer"),
    ]
    print(f"\n{BOLD}Available actions:{RESET}")
    for num, at, _, desc in actions:
        print(f"  {CYAN}{num}{RESET}  {at:25s}  {DIM}{desc}{RESET}")
    print()


def prompt_action(env) -> Action:
    """Prompt the user to choose and configure an action interactively."""
    action_map = {
        "1": ActionType.LIST_FUNCTIONS,
        "2": ActionType.GET_FUNCTION_CODE,
        "3": ActionType.GET_FUNCTION_SUMMARY,
        "4": ActionType.GET_FILE_METADATA,
        "5": ActionType.GET_STATE_VARIABLE,
        "6": ActionType.GET_CALL_GRAPH,
        "7": ActionType.SUBMIT,
    }
    while True:
        choice = input(f"{BOLD}Choose action (1-7): {RESET}").strip()
        if choice not in action_map:
            print(f"  {YELLOW}Enter a number 1–7{RESET}")
            continue
        at = action_map[choice]
        params = {}

        if at in (ActionType.GET_FUNCTION_CODE, ActionType.GET_FUNCTION_SUMMARY):
            fn = input("  Function name: ").strip()
            params = {"function_name": fn}

        elif at == ActionType.GET_STATE_VARIABLE:
            var = input("  Variable name (leave blank to list all): ").strip()
            if var:
                params = {"variable_name": var}

        elif at == ActionType.SUBMIT:
            fn   = input("  Vulnerable function name: ").strip()
            vuln = input("  Vulnerability type (2-3 words): ").strip()
            params = {"function_name": fn, "vulnerability_type": vuln}

        return Action(action_type=at, params=params)


# ─────────────────────────────────────────────────────────────────────────────
# Scripted demo agent
# ─────────────────────────────────────────────────────────────────────────────

DEMO_SCRIPTS = {
    # seed → list of (ActionType, params, commentary)
    42: [
        (ActionType.GET_FILE_METADATA, {},
         "First, get high-level contract info to understand the domain."),
        (ActionType.LIST_FUNCTIONS, {},
         "List functions to understand the attack surface."),
        (ActionType.GET_FUNCTION_SUMMARY, {"function_name": "emergencyDrain"},
         "emergencyDrain sounds dangerous — check what it's supposed to do."),
        (ActionType.GET_FUNCTION_CODE, {"function_name": "emergencyDrain"},
         "Inspect the code — no onlyOwner modifier! Anyone can drain the vault."),
        (ActionType.SUBMIT, {"function_name": "emergencyDrain", "vulnerability_type": "missing access control"},
         "Confident: missing access control. Submitting!"),
    ],
    7: [
        (ActionType.LIST_FUNCTIONS, {},
         "Start by surveying all functions."),
        (ActionType.GET_FUNCTION_SUMMARY, {"function_name": "finalize"},
         "finalize — what does this auction close-out function do?"),
        (ActionType.GET_FUNCTION_CODE, {"function_name": "finalize"},
         "Uses block.timestamp for deadline check — miners can manipulate this."),
        (ActionType.SUBMIT, {"function_name": "finalize", "vulnerability_type": "timestamp dependence"},
         "Classic timestamp manipulation. Submitting."),
    ],
}

DEFAULT_DEMO_SEED = 42


def run_auto_demo(seed: int, delay: float = 0.9):
    """Run the scripted demo agent with printed commentary."""
    script = DEMO_SCRIPTS.get(seed)
    if script is None:
        # Generic fallback: list, code of first suspicious fn, submit
        print(f"{YELLOW}No pre-written script for seed {seed}. Running generic agent.{RESET}\n")
        script = [
            (ActionType.LIST_FUNCTIONS, {}, "Listing all functions first."),
            (ActionType.GET_FILE_METADATA, {}, "Checking contract metadata."),
        ]

    env = Task1Environment()
    result = env.reset(seed=seed)
    obs = result.observation

    banner()
    print(f"{BOLD}Mode:{RESET} Automated demo  |  {BOLD}Seed:{RESET} {seed}\n")
    print_observation(obs)

    for at, params, commentary in script:
        time.sleep(delay)
        print(f"\n{CYAN}▶ Agent thinking:{RESET} {commentary}")
        time.sleep(delay * 0.5)
        action = Action(action_type=at, params=params)
        step   = env.step(action)
        obs    = step.observation
        print_observation(obs)

        if step.done:
            _print_episode_summary(obs)
            return

    # Episode not finished — shouldn't happen with a good script
    state = env.state()
    print(f"\n{YELLOW}Episode not completed (no submit action in script). "
          f"Target was: {state.target_function}{RESET}")


# ─────────────────────────────────────────────────────────────────────────────
# Interactive mode
# ─────────────────────────────────────────────────────────────────────────────

def run_interactive(seed: int = None):
    env  = Task1Environment()
    seed = seed or 42
    result = env.reset(seed=seed)
    obs    = result.observation

    banner()
    print(f"{BOLD}Mode:{RESET} Interactive  |  {BOLD}Seed:{RESET} {seed}")
    print(f"{DIM}Tip: Start with list_functions and get_file_metadata.{RESET}\n")
    print_observation(obs)

    while not obs.done:
        print_action_menu()
        try:
            action = prompt_action(env)
        except (KeyboardInterrupt, EOFError):
            print(f"\n{YELLOW}Demo interrupted.{RESET}")
            break

        step = env.step(action)
        obs  = step.observation
        print()
        print_observation(obs)

        if step.done:
            _print_episode_summary(obs)
            break

    _offer_replay()


def _print_episode_summary(obs):
    print()
    print(f"{BOLD}{'═' * 64}{RESET}")
    reward = obs.cumulative_reward
    colour = GREEN if reward > 0 else RED
    print(f"{BOLD}Episode complete!{RESET}")
    print(f"  Steps taken   : {obs.step_count}")
    print(f"  Total reward  : {colour}{reward:+.2f}{RESET}")
    last = obs.last_action_result or ""
    if "✅ CORRECT" in last or "EXCELLENT" in last:
        print(f"  {GREEN}Perfect score — full marks!{RESET}")
    elif "⚠️" in last or "PARTIAL" in last:
        print(f"  {YELLOW}Partial credit.{RESET}")
    elif "🟡 GOOD" in last:
        print(f"  {YELLOW}Good — most key concepts matched!{RESET}")
    elif "🟠" in last:
        print(f"  {YELLOW}Partial — some key concepts matched.{RESET}")
    elif "❌" in last:
        print(f"  {RED}Incorrect — better luck next episode.{RESET}")
    else:
        print(f"  {'Good effort!' if reward > 0 else 'Keep exploring next time.'}")
    print(f"{BOLD}{'═' * 64}{RESET}\n")


def _offer_replay():
    try:
        again = input("Play again? (y/n): ").strip().lower()
        if again == "y":
            run_interactive()
    except (KeyboardInterrupt, EOFError):
        pass


# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Smart Contract Audit RL Environment — Demo"
    )
    parser.add_argument(
        "--auto", action="store_true",
        help="Run the scripted demo agent (no keyboard input required)"
    )
    parser.add_argument(
        "--seed", type=int, default=DEFAULT_DEMO_SEED,
        help="Episode seed (default: 42)"
    )
    parser.add_argument(
        "--delay", type=float, default=0.9,
        help="Seconds between auto-agent steps (default: 0.9)"
    )
    args = parser.parse_args()

    if args.auto:
        run_auto_demo(seed=args.seed, delay=args.delay)
    else:
        run_interactive(seed=args.seed)


if __name__ == "__main__":
    main()


# ─────────────────────────────────────────────────────────────────────────────
# Task 2 demo
# ─────────────────────────────────────────────────────────────────────────────

DEMO_SCRIPTS_T2 = {
    42: [
        (ActionType.GET_FUNCTION_NATSPEC,  {}, "First, read the NatSpec to understand intent and expected outputs."),
        (ActionType.GET_IO,                {}, "Check parameters, return type and expected behaviour."),
        (ActionType.GET_FUNCTION_CODE,     {}, "Read the actual Solidity code to confirm the behaviour."),
        (ActionType.SUBMIT_PROPERTY,
         {"property": "After a successful claimRewards call, all of the caller's accrued reward tokens are transferred to the caller and their rewards balance is set to zero. Reverts if the caller has no accrued rewards."},
         "Confident about the property. Submitting!"),
    ],
}


def run_auto_demo_t2(seed: int = 42, delay: float = 0.9):
    """Run the scripted Task 2 demo."""
    from tasks.task2.environment import Task2Environment

    script = DEMO_SCRIPTS_T2.get(seed)
    env = Task2Environment()
    result = env.reset(seed=seed)
    obs = result.observation

    print()
    print(f"{BOLD}{CYAN}╔══════════════════════════════════════════════════════════╗")
    print(f"║   Smart Contract Audit RL Env  ·  Task 2 Demo            ║")
    print(f"╚══════════════════════════════════════════════════════════╝{RESET}")
    print()
    print(f"{BOLD}Mode:{RESET} Automated demo  |  {BOLD}Seed:{RESET} {seed}")
    print(f"{BOLD}Task:{RESET} Property Discovery")
    print()

    fn_name = obs.extra.get("target_function", "?")
    sig     = obs.extra.get("target_signature", "")
    print(f"{BOLD}Contract :{RESET} {obs.contract_name}")
    print(f"{BOLD}Function :{RESET} {fn_name}  ({sig})")
    print(f"{BOLD}Goal     :{RESET} Write the natural-language property for '{fn_name}'")
    print(DIVIDER)

    if not script:
        print(f"{YELLOW}No pre-written script for seed {seed}.{RESET}")
        return

    for at, params, commentary in script:
        time.sleep(delay)
        print(f"\n{CYAN}▶ Agent thinking:{RESET} {commentary}")
        time.sleep(delay * 0.5)
        step_result = env.step(Action(action_type=at, params=params))
        sobs = step_result.observation
        print(DIVIDER)
        print(f"{BOLD}Step {sobs.step_count:2d}{RESET}  [{at.value}]  r={step_result.reward.value:+.2f}  cum={sobs.cumulative_reward:+.2f}")
        result_text = sobs.last_action_result or ""
        colour = GREEN if step_result.reward.value > 0 else (YELLOW if step_result.reward.value > -0.15 else YELLOW)
        for line in result_text.split("\n")[:8]:
            print(f"  {colour}{line[:90]}{RESET}")
        print(DIVIDER)

        if step_result.done:
            _print_episode_summary(sobs)
            return


# ─────────────────────────────────────────────────────────────────────────────
# Task 3 demo
# ─────────────────────────────────────────────────────────────────────────────

DEMO_SCRIPTS_T3 = {
    42: [
        (ActionType.GET_PROPERTY_SPECIFICATION, {},
         "Read the formal spec first — cheapest action at -0.03."),
        (ActionType.LIST_FUNCTIONS, {},
         "List all functions to survey candidates."),
        (ActionType.GET_FUNCTION_CODE, {"function_name": "emergencyDrain"},
         "No access modifier! Anyone can call this — that's the violation."),
        (ActionType.SUBMIT_FUNCTION, {"function_name": "emergencyDrain"},
         "Confident. emergencyDrain violates the access-control property."),
    ],
    45: [
        (ActionType.GET_PROPERTY_SPECIFICATION, {},
         "Formal spec: first caller at valid price should win."),
        (ActionType.LIST_FUNCTIONS, {},
         "Auction contract — bid() immediately looks suspicious."),
        (ActionType.GET_FUNCTION_CODE, {"function_name": "bid"},
         "No commit-reveal, no maxPrice guard — front-running is trivially possible."),
        (ActionType.SUBMIT_FUNCTION, {"function_name": "bid"},
         "bid() violates the front-running property. Submitting."),
    ],
}


def run_auto_demo_t3(seed: int = 42, delay: float = 0.9):
    """Run the scripted Task 3 demo."""
    from tasks.task3.environment import Task3Environment

    script = DEMO_SCRIPTS_T3.get(seed)
    env    = Task3Environment()
    result = env.reset(seed=seed)
    obs    = result.observation

    print()
    print(f"{BOLD}{CYAN}╔══════════════════════════════════════════════════════════╗")
    print(f"║   Smart Contract Audit RL Env  ·  Task 3 Demo            ║")
    print(f"╚══════════════════════════════════════════════════════════╝{RESET}")
    print()
    print(f"{BOLD}Mode:{RESET} Automated demo  |  {BOLD}Seed:{RESET} {seed}")
    print(f"{BOLD}Task:{RESET} Rule Checker")
    print()

    prop = obs.extra.get("property_english", "")
    print(f"{BOLD}Contract  :{RESET} {obs.contract_name}")
    print(f"{BOLD}Property  :{RESET} {prop[:100]}{'...' if len(prop) > 100 else ''}")
    print(f"{BOLD}Goal      :{RESET} Find the function that violates this property.")
    print(DIVIDER)

    if not script:
        print(f"{YELLOW}No pre-written script for seed {seed}. Try seed 42 or 45.{RESET}")
        return

    for at, params, commentary in script:
        time.sleep(delay)
        print(f"\n{CYAN}▶ Agent thinking:{RESET} {commentary}")
        time.sleep(delay * 0.5)
        step_result = env.step(Action(action_type=at, params=params))
        sobs = step_result.observation
        print(DIVIDER)
        print(f"{BOLD}Step {sobs.step_count:2d}{RESET}  [{at.value}]  "
              f"r={step_result.reward.value:+.2f}  cum={sobs.cumulative_reward:+.2f}")
        result_text = sobs.last_action_result or ""
        colour = GREEN if step_result.reward.value > 0 else YELLOW
        for line in result_text.split("\n")[:6]:
            print(f"  {colour}{line[:90]}{RESET}")
        print(DIVIDER)

        if step_result.done:
            _print_episode_summary(sobs)
            return