reset.py

#!/usr/bin/env python3
"""
reset.py - Clean slate for debugging cycles.

Removes checkpoints, training data, and memory files to start fresh.
Use with caution - this is destructive!

Usage:
    python reset.py              # Interactive mode (asks for confirmation)
    python reset.py --force      # Skip confirmation (for scripting)
    python reset.py --dry-run    # Show what would be deleted
    python reset.py --keep-logs  # Keep ALL log files (*_sessions.jsonl, *_graded.jsonl,
                                 #   *_meta.jsonl, *_summary.json, train_*.log, *_transcript.txt)
                                 #   This preserves training data so you can retrain from scratch.
    python reset.py --reset-model  # Also delete and recreate basil_v001
"""

import os
import sys
import shutil
import argparse
from datetime import datetime

from config import (
    BASE_DIR, MODELS_DIR, LOG_DIR, MEMORY_DIR, CHECKPOINT_DIR,
    SESSION_SUMMARIES_DIR, METRICS_FILE, ROTATION_STATE_FILE,
    BASIL_ASSESSMENT_FILE,
    SESSION_METRICS_DIR,
    SUBJECTS_FILE, USED_LESSONS_FILE, USED_STORIES_FILE, USED_TOPICS_FILE,
    STORY_GENRES_FILE, TEACHING_ANGLES_FILE, SCIENCE_DOMAINS_FILE,
    SUBJECT_TOPICS_DIR,
    TASK_NATURALIZER_CACHE_FILE,
    IDENTITY_FILE,
)

# Model to preserve (initial random weights - expensive to recreate)
PRESERVED_MODEL = "basil_v001"


# ANSI colors for terminal output
class Colors:
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    BOLD = '\033[1m'
    END = '\033[0m'


def sizeof_fmt(num, suffix='B'):
    """Human-readable file size."""
    for unit in ['', 'K', 'M', 'G', 'T']:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}P{suffix}"


def get_dir_size(path):
    """Get total size of a directory."""
    total = 0
    if os.path.exists(path):
        for dirpath, dirnames, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                if os.path.exists(fp):
                    total += os.path.getsize(fp)
    return total


def list_models(include_preserved=False):
    """List all basil_v* model directories except the initial random weights model.
    
    If include_preserved=True, also includes basil_v001 (for --reset-model).
    """
    if not os.path.exists(MODELS_DIR):
        return []
    return sorted([
        d for d in os.listdir(MODELS_DIR)
        if d.startswith("basil_v") and os.path.isdir(os.path.join(MODELS_DIR, d))
        and (include_preserved or d != PRESERVED_MODEL)  # Keep initial random weights model unless told otherwise
    ])


def list_logs(include_legacy=True):
    """List all training log files."""
    if not os.path.exists(LOG_DIR):
        return [], []
    
    graded_logs = []
    legacy_logs = []
    
    for fname in os.listdir(LOG_DIR):
        fpath = os.path.join(LOG_DIR, fname)
        if not os.path.isfile(fpath):
            continue
        if fname.endswith("_graded.jsonl") or fname.endswith("_structured.jsonl"):
            graded_logs.append(fname)
        elif fname.endswith(".jsonl") and include_legacy:
            legacy_logs.append(fname)
    
    return sorted(graded_logs), sorted(legacy_logs)


def list_memory_files():
    """List all memory/state files."""
    files = []
    for fpath in [METRICS_FILE, ROTATION_STATE_FILE, BASIL_ASSESSMENT_FILE]:
        if os.path.exists(fpath):
            files.append(fpath)
    return files


def list_session_summaries():
    """List all session summary files."""
    if not os.path.exists(SESSION_SUMMARIES_DIR):
        return []
    return [
        os.path.join(SESSION_SUMMARIES_DIR, f)
        for f in os.listdir(SESSION_SUMMARIES_DIR)
        if os.path.isfile(os.path.join(SESSION_SUMMARIES_DIR, f))
    ]


def list_session_metrics():
    """List all session metrics files."""
    if not os.path.exists(SESSION_METRICS_DIR):
        return []
    return [
        os.path.join(SESSION_METRICS_DIR, f)
        for f in os.listdir(SESSION_METRICS_DIR)
        if os.path.isfile(os.path.join(SESSION_METRICS_DIR, f))
    ]


def list_checkpoints():
    """List checkpoint directories."""
    if not os.path.exists(CHECKPOINT_DIR):
        return []
    return [
        d for d in os.listdir(CHECKPOINT_DIR)
        if os.path.isdir(os.path.join(CHECKPOINT_DIR, d))
    ]


def list_classroom_files():
    """List classroom runtime state files (individual files + subject_topics dir)."""
    files = []
    for fpath in [SUBJECTS_FILE, USED_LESSONS_FILE, USED_STORIES_FILE, USED_TOPICS_FILE, STORY_GENRES_FILE, TEACHING_ANGLES_FILE, SCIENCE_DOMAINS_FILE]:
        if os.path.exists(fpath):
            files.append(fpath)
    # Per-subject topic cache directory
    if os.path.isdir(SUBJECT_TOPICS_DIR):
        for fname in os.listdir(SUBJECT_TOPICS_DIR):
            fpath = os.path.join(SUBJECT_TOPICS_DIR, fname)
            if os.path.isfile(fpath):
                files.append(fpath)
    return files


def list_training_logs():
    """List training log files (train_*.log)."""
    if not os.path.exists(LOG_DIR):
        return []
    return [
        os.path.join(LOG_DIR, f)
        for f in os.listdir(LOG_DIR)
        if f.startswith("train_") and f.endswith(".log")
    ]


def list_transcript_files():
    """List session transcript files (*_transcript.txt)."""
    if not os.path.exists(LOG_DIR):
        return []
    return [
        os.path.join(LOG_DIR, f)
        for f in os.listdir(LOG_DIR)
        if f.endswith("_transcript.txt")
    ]


def list_batch_files():
    """List batch-related files."""
    if not os.path.exists(LOG_DIR):
        return []
    batch_files = []
    for f in os.listdir(LOG_DIR):
        if f.startswith("batch_") and (
            f.endswith("_summary.json") or
            f.endswith("_meta.jsonl") or
            f.endswith("_sessions.jsonl")
        ):
            batch_files.append(os.path.join(LOG_DIR, f))
    return sorted(batch_files)


def list_identity_files():
    """List identity log files."""
    files = []
    if os.path.exists(IDENTITY_FILE):
        files.append(IDENTITY_FILE)
    return files


def list_cache_files():
    """List cache files."""
    files = []
    if os.path.exists(TASK_NATURALIZER_CACHE_FILE):
        files.append(TASK_NATURALIZER_CACHE_FILE)
    return files


def print_summary(keep_logs=False, reset_model=False):
    """Print summary of what will be deleted."""
    print(f"\n{Colors.BOLD}=== Reset Summary ==={Colors.END}\n")
    
    total_size = 0
    total_items = 0
    
    # Models (basil_v001 is preserved unless --reset-model)
    models = list_models(include_preserved=reset_model)
    if models:
        model_size = sum(get_dir_size(os.path.join(MODELS_DIR, m)) for m in models)
        total_size += model_size
        total_items += len(models)
        print(f"{Colors.RED}Models ({len(models)} dirs, {sizeof_fmt(model_size)}):{Colors.END}")
        for m in models:
            if m == PRESERVED_MODEL:
                print(f"  - {m} {Colors.YELLOW}(will be deleted and recreated){Colors.END}")
            else:
                print(f"  - {m}")
        if not reset_model:
            print(f"{Colors.GREEN}  (keeping {PRESERVED_MODEL} - initial random weights){Colors.END}")
    else:
        print(f"{Colors.GREEN}Models: None to delete (keeping {PRESERVED_MODEL}){Colors.END}")
    
    # All log-related files: graded logs, training logs, transcripts, batch files
    if keep_logs:
        # Count what we're keeping for display
        graded_logs, legacy_logs = list_logs(include_legacy=True)
        training_logs = list_training_logs()
        transcript_files = list_transcript_files()
        batch_files = list_batch_files()
        kept_count = len(graded_logs) + len(legacy_logs) + len(training_logs) + len(transcript_files) + len(batch_files)
        if kept_count > 0:
            print(f"\n{Colors.GREEN}Logs: KEEPING all {kept_count} files in logs/ (--keep-logs){Colors.END}")
            if batch_files:
                sessions = [f for f in batch_files if f.endswith("_sessions.jsonl")]
                print(f"  {Colors.GREEN}Includes {len(sessions)} session files (training data){Colors.END}")
        else:
            print(f"\n{Colors.GREEN}Logs: None found{Colors.END}")
    else:
        # Not keeping logs — show what will be deleted
        graded_logs, legacy_logs = list_logs(include_legacy=True)
        if graded_logs or legacy_logs:
            log_size = sum(
                os.path.getsize(os.path.join(LOG_DIR, f))
                for f in graded_logs + legacy_logs
                if os.path.exists(os.path.join(LOG_DIR, f))
            )
            total_size += log_size
            total_items += len(graded_logs) + len(legacy_logs)
            
            print(f"\n{Colors.RED}Logs ({len(graded_logs) + len(legacy_logs)} files, {sizeof_fmt(log_size)}):{Colors.END}")
            if graded_logs:
                print(f"  Graded: {len(graded_logs)} files")
                for f in graded_logs[:3]:
                    print(f"    - {f}")
                if len(graded_logs) > 3:
                    print(f"    ... and {len(graded_logs) - 3} more")
            if legacy_logs:
                print(f"  Legacy: {len(legacy_logs)} files")
                for f in legacy_logs[:3]:
                    print(f"    - {f}")
                if len(legacy_logs) > 3:
                    print(f"    ... and {len(legacy_logs) - 3} more")
        else:
            print(f"\n{Colors.GREEN}Logs: None to delete{Colors.END}")
        
        # Training logs
        training_logs = list_training_logs()
        if training_logs:
            train_log_size = sum(os.path.getsize(f) for f in training_logs if os.path.exists(f))
            total_size += train_log_size
            total_items += len(training_logs)
            print(f"\n{Colors.RED}Training Logs ({len(training_logs)} files, {sizeof_fmt(train_log_size)}):{Colors.END}")
            for f in training_logs[:3]:
                print(f"  - {os.path.basename(f)}")
            if len(training_logs) > 3:
                print(f"  ... and {len(training_logs) - 3} more")
        else:
            print(f"\n{Colors.GREEN}Training Logs: None found{Colors.END}")
        
        # Transcript files
        transcript_files = list_transcript_files()
        if transcript_files:
            transcript_size = sum(os.path.getsize(f) for f in transcript_files if os.path.exists(f))
            total_size += transcript_size
            total_items += len(transcript_files)
            print(f"\n{Colors.RED}Transcript Files ({len(transcript_files)} files, {sizeof_fmt(transcript_size)}):{Colors.END}")
            for f in transcript_files[:3]:
                print(f"  - {os.path.basename(f)}")
            if len(transcript_files) > 3:
                print(f"  ... and {len(transcript_files) - 3} more")
        else:
            print(f"\n{Colors.GREEN}Transcript Files: None found{Colors.END}")
        
        # Batch files
        batch_files = list_batch_files()
        if batch_files:
            batch_size = sum(os.path.getsize(f) for f in batch_files if os.path.exists(f))
            total_size += batch_size
            total_items += len(batch_files)
            print(f"\n{Colors.RED}Batch Files ({len(batch_files)} files, {sizeof_fmt(batch_size)}):{Colors.END}")
            for f in batch_files[:3]:
                print(f"  - {os.path.basename(f)}")
            if len(batch_files) > 3:
                print(f"  ... and {len(batch_files) - 3} more")
        else:
            print(f"\n{Colors.GREEN}Batch Files: None found{Colors.END}")
    
    # Classroom files
    classroom_files = list_classroom_files()
    if classroom_files:
        classroom_size = sum(os.path.getsize(f) for f in classroom_files if os.path.exists(f))
        total_size += classroom_size
        total_items += len(classroom_files)
        print(f"\n{Colors.RED}Classroom Files ({len(classroom_files)} files, {sizeof_fmt(classroom_size)}):{Colors.END}")
        for f in classroom_files:
            print(f"  - {os.path.basename(f)}")
    else:
        print(f"\n{Colors.GREEN}Classroom Files: None found{Colors.END}")
    
    # Cache files
    cache_files = list_cache_files()
    if cache_files:
        cache_size = sum(os.path.getsize(f) for f in cache_files if os.path.exists(f))
        total_size += cache_size
        total_items += len(cache_files)
        print(f"\n{Colors.RED}Cache Files ({len(cache_files)} files, {sizeof_fmt(cache_size)}):{Colors.END}")
        for f in cache_files:
            print(f"  - {os.path.basename(f)}")
    else:
        print(f"\n{Colors.GREEN}Cache Files: None found{Colors.END}")
    
    # Identity files
    identity_files = list_identity_files()
    if identity_files:
        identity_size = sum(os.path.getsize(f) for f in identity_files if os.path.exists(f))
        total_size += identity_size
        total_items += len(identity_files)
        print(f"\n{Colors.RED}Identity Files ({len(identity_files)} files, {sizeof_fmt(identity_size)}):{Colors.END}")
        for f in identity_files:
            print(f"  - {os.path.basename(f)}")
    else:
        print(f"\n{Colors.GREEN}Identity Files: None found{Colors.END}")
    
    # Memory files
    memory_files = list_memory_files()
    session_summaries = list_session_summaries()
    session_metrics = list_session_metrics()
    all_memory = memory_files + session_summaries + session_metrics
    if all_memory:
        mem_size = sum(os.path.getsize(f) for f in all_memory if os.path.exists(f))
        total_size += mem_size
        total_items += len(all_memory)
        
        print(f"\n{Colors.RED}Memory ({len(all_memory)} files, {sizeof_fmt(mem_size)}):{Colors.END}")
        for f in memory_files:
            print(f"  - {os.path.basename(f)}")
        if session_summaries:
            print(f"  + {len(session_summaries)} session summaries")
        if session_metrics:
            print(f"  + {len(session_metrics)} session metrics")
    else:
        print(f"\n{Colors.GREEN}Memory: None found{Colors.END}")
    
    # Checkpoints
    checkpoints = list_checkpoints()
    if checkpoints:
        ckpt_size = sum(get_dir_size(os.path.join(CHECKPOINT_DIR, c)) for c in checkpoints)
        total_size += ckpt_size
        total_items += len(checkpoints)
        print(f"\n{Colors.RED}Checkpoints ({len(checkpoints)} dirs, {sizeof_fmt(ckpt_size)}):{Colors.END}")
        for c in checkpoints[:3]:
            print(f"  - {c}")
        if len(checkpoints) > 3:
            print(f"  ... and {len(checkpoints) - 3} more")
    else:
        print(f"\n{Colors.GREEN}Checkpoints: None found{Colors.END}")
    
    print(f"\n{Colors.BOLD}Total: {total_items} items, {sizeof_fmt(total_size)}{Colors.END}")
    
    return total_items > 0


def do_reset(dry_run=False, keep_logs=False, reset_model=False):
    """Perform the actual reset."""
    deleted_count = 0
    
    # Delete models
    models = list_models(include_preserved=reset_model)
    for m in models:
        mpath = os.path.join(MODELS_DIR, m)
        if dry_run:
            print(f"[DRY RUN] Would delete: {mpath}")
        else:
            shutil.rmtree(mpath)
            print(f"Deleted: {mpath}")
        deleted_count += 1
    
    # Delete logs, training logs, transcripts, batch files
    # --keep-logs protects ALL files in the logs/ directory
    if keep_logs:
        print(f"Keeping all log files (--keep-logs)")
    else:
        # Delete graded and legacy logs
        graded_logs, legacy_logs = list_logs(include_legacy=True)
        for f in graded_logs + legacy_logs:
            fpath = os.path.join(LOG_DIR, f)
            if dry_run:
                print(f"[DRY RUN] Would delete: {fpath}")
            else:
                os.remove(fpath)
                print(f"Deleted: {fpath}")
            deleted_count += 1
        
        # Delete training logs
        training_logs = list_training_logs()
        for f in training_logs:
            if dry_run:
                print(f"[DRY RUN] Would delete: {f}")
            else:
                os.remove(f)
                print(f"Deleted: {f}")
            deleted_count += 1
        
        # Delete transcript files
        transcript_files = list_transcript_files()
        for f in transcript_files:
            if dry_run:
                print(f"[DRY RUN] Would delete: {f}")
            else:
                os.remove(f)
                print(f"Deleted: {f}")
            deleted_count += 1
        
        # Delete batch files (includes *_sessions.jsonl — the training data!)
        batch_files = list_batch_files()
        for f in batch_files:
            if dry_run:
                print(f"[DRY RUN] Would delete: {f}")
            else:
                os.remove(f)
                print(f"Deleted: {f}")
            deleted_count += 1
    
    # Delete memory files
    memory_files = list_memory_files()
    for f in memory_files:
        if dry_run:
            print(f"[DRY RUN] Would delete: {f}")
        else:
            os.remove(f)
            print(f"Deleted: {f}")
        deleted_count += 1
    
    # Delete session summaries
    session_summaries = list_session_summaries()
    for f in session_summaries:
        if dry_run:
            print(f"[DRY RUN] Would delete: {f}")
        else:
            os.remove(f)
            print(f"Deleted: {f}")
        deleted_count += 1
    
    # Delete session metrics
    session_metrics = list_session_metrics()
    for f in session_metrics:
        if dry_run:
            print(f"[DRY RUN] Would delete: {f}")
        else:
            os.remove(f)
            print(f"Deleted: {f}")
        deleted_count += 1
    
    # Delete checkpoints
    checkpoints = list_checkpoints()
    for c in checkpoints:
        cpath = os.path.join(CHECKPOINT_DIR, c)
        if dry_run:
            print(f"[DRY RUN] Would delete: {cpath}")
        else:
            shutil.rmtree(cpath)
            print(f"Deleted: {cpath}")
        deleted_count += 1
    
    # Delete classroom files
    classroom_files = list_classroom_files()
    for f in classroom_files:
        if dry_run:
            print(f"[DRY RUN] Would delete: {f}")
        else:
            os.remove(f)
            print(f"Deleted: {f}")
        deleted_count += 1
    
    # Delete cache files
    cache_files = list_cache_files()
    for f in cache_files:
        if dry_run:
            print(f"[DRY RUN] Would delete: {f}")
        else:
            os.remove(f)
            print(f"Deleted: {f}")
        deleted_count += 1
    
    # Delete identity files
    identity_files = list_identity_files()
    for f in identity_files:
        if dry_run:
            print(f"[DRY RUN] Would delete: {f}")
        else:
            os.remove(f)
            print(f"Deleted: {f}")
        deleted_count += 1
    
    return deleted_count


def create_fresh_basil(auto=False):
    """Offer to create a fresh basil_v001 if it doesn't exist.
    
    If auto=True, skip the interactive prompt and just create it.
    """
    v001_path = os.path.join(MODELS_DIR, PRESERVED_MODEL)
    if os.path.exists(v001_path):
        print(f"\n{Colors.GREEN}✓ {PRESERVED_MODEL} preserved (initial random weights){Colors.END}")
        return
    
    if not auto:
        print(f"\n{Colors.YELLOW}{PRESERVED_MODEL} not found. Create it?{Colors.END}")
        response = input("Create initial model? [y/N]: ").strip().lower()
        if response != 'y':
            return
    else:
        print(f"\n{Colors.BLUE}Creating fresh {PRESERVED_MODEL}...{Colors.END}")
    
    try:
        # Import here to avoid circular dependencies
        import subprocess
        result = subprocess.run(
            [sys.executable, os.path.join(BASE_DIR, "create_basil_v0001.py")],
            capture_output=True,
            text=True
        )
        if result.returncode == 0:
            print(f"{Colors.GREEN}✓ Fresh {PRESERVED_MODEL} created{Colors.END}")
        else:
            print(f"{Colors.RED}Failed to create model: {result.stderr}{Colors.END}")
    except Exception as e:
        print(f"{Colors.RED}Error creating model: {e}{Colors.END}")


def main():
    parser = argparse.ArgumentParser(
        description="Reset Bootstrap Basil to a clean state for debugging.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python reset.py              # Interactive mode
  python reset.py --dry-run    # Preview what would be deleted
  python reset.py --force      # Delete without confirmation
  python reset.py --keep-logs  # Keep ALL log files (sessions, graded, training logs, transcripts)
  python reset.py --force --reset-model --keep-logs  # Nuke model, keep training data for retrain
  python reset.py --force --reset-model               # Full nuke including training data
"""
    )
    parser.add_argument(
        "--force", "-f",
        action="store_true",
        help="Skip confirmation prompt"
    )
    parser.add_argument(
        "--dry-run", "-n",
        action="store_true",
        help="Show what would be deleted without actually deleting"
    )
    parser.add_argument(
        "--keep-logs",
        action="store_true",
        help="Keep ALL files in logs/ directory (sessions, graded, training logs, transcripts, batch files). Use this to preserve training data when resetting model weights."
    )
    parser.add_argument(
        "--create-fresh",
        action="store_true",
        help="Create fresh basil_v001 after reset"
    )
    parser.add_argument(
        "--reset-model",
        action="store_true",
        help="Also delete and recreate the initial random weights model (basil_v001)"
    )
    args = parser.parse_args()
    
    print(f"\n{Colors.BOLD}{Colors.RED}⚠️  Bootstrap Basil Reset Tool{Colors.END}")
    print(f"This will delete training data and checkpoints.\n")
    
    # Show summary
    has_items = print_summary(keep_logs=args.keep_logs, reset_model=args.reset_model)
    
    if not has_items:
        print(f"\n{Colors.GREEN}Nothing to delete. Already clean!{Colors.END}\n")
        return
    
    if args.dry_run:
        print(f"\n{Colors.YELLOW}=== DRY RUN ==={Colors.END}\n")
        do_reset(dry_run=True, keep_logs=args.keep_logs, reset_model=args.reset_model)
        print(f"\n{Colors.YELLOW}No files were deleted (dry run).{Colors.END}\n")
        return
    
    # Confirmation
    if not args.force:
        print(f"\n{Colors.RED}{Colors.BOLD}This action cannot be undone!{Colors.END}")
        response = input(f"Type 'RESET' to confirm: ").strip()
        if response != "RESET":
            print(f"\n{Colors.YELLOW}Aborted.{Colors.END}\n")
            return
    
    # Do the reset
    print(f"\n{Colors.BOLD}Resetting...{Colors.END}\n")
    deleted = do_reset(dry_run=False, keep_logs=args.keep_logs, reset_model=args.reset_model)
    
    print(f"\n{Colors.GREEN}✓ Reset complete. Deleted {deleted} items.{Colors.END}")
    
    # Offer to create fresh model
    # --reset-model implies recreating it; --create-fresh also triggers this
    if args.reset_model or args.create_fresh:
        create_fresh_basil(auto=True)
    elif not args.force and not list_models():
        create_fresh_basil(auto=False)
    
    print()


if __name__ == "__main__":
    main()