#!/usr/bin/env python3
"""
Context Compaction Manager

Handles pre-compact state saving and post-compact resume injection.
Monitors context usage and triggers appropriate hooks.

Usage:
  python3 context_compact.py save [--workflow-dir .workflow/versions/v001]
  python3 context_compact.py resume [--workflow-dir .workflow/versions/v001]
  python3 context_compact.py status [--workflow-dir .workflow/versions/v001]
"""

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional


# ============================================================================
# Configuration
# ============================================================================

DEFAULT_WORKFLOW_DIR = ".workflow/versions/v001"
STATE_FILE = "context_state.json"
RESUME_PROMPT_FILE = "resume_prompt.md"
MODIFIED_FILES_FILE = "modified_files.json"

# Context thresholds (percentage)
THRESHOLDS = {
    "warning": 0.70,
    "save": 0.80,
    "compact": 0.90,
    "critical": 0.95
}


# ============================================================================
# File Operations
# ============================================================================

def load_json(filepath: str) -> dict:
    """Load JSON file."""
    if not os.path.exists(filepath):
        return {}
    try:
        with open(filepath, 'r') as f:
            return json.load(f)
    except json.JSONDecodeError:
        return {}


def save_json(filepath: str, data: dict):
    """Save data to JSON file."""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, 'w') as f:
        json.dump(data, f, indent=2, default=str)


def save_text(filepath: str, content: str):
    """Save text to file."""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, 'w') as f:
        f.write(content)


# ============================================================================
# Git Operations
# ============================================================================

def get_git_status() -> List[Dict[str, str]]:
    """Get list of modified files from git."""
    try:
        result = subprocess.run(
            ['git', 'status', '--porcelain'],
            capture_output=True, text=True, check=True
        )
        files = []
        for line in result.stdout.strip().split('\n'):
            if line:
                status = line[:2].strip()
                path = line[3:]
                action = {
                    'M': 'modified',
                    'A': 'added',
                    'D': 'deleted',
                    '?': 'untracked',
                    'R': 'renamed'
                }.get(status[0] if status else '?', 'unknown')
                files.append({'path': path, 'action': action, 'summary': ''})
        return files
    except subprocess.CalledProcessError:
        return []


def get_recent_commits(count: int = 5) -> List[Dict[str, str]]:
    """Get recent commit messages."""
    try:
        result = subprocess.run(
            ['git', 'log', f'-{count}', '--oneline'],
            capture_output=True, text=True, check=True
        )
        commits = []
        for line in result.stdout.strip().split('\n'):
            if line:
                parts = line.split(' ', 1)
                commits.append({
                    'hash': parts[0],
                    'message': parts[1] if len(parts) > 1 else ''
                })
        return commits
    except subprocess.CalledProcessError:
        return []


def create_checkpoint(message: str = "WIP: Pre-compaction checkpoint"):
    """Create a git checkpoint with uncommitted changes."""
    try:
        # Check if there are changes
        result = subprocess.run(
            ['git', 'status', '--porcelain'],
            capture_output=True, text=True
        )
        if result.stdout.strip():
            # Stage all changes
            subprocess.run(['git', 'add', '-A'], check=True)
            # Commit
            subprocess.run(['git', 'commit', '-m', message], check=True)
            print(f"Created checkpoint: {message}")
            return True
    except subprocess.CalledProcessError as e:
        print(f"Warning: Could not create checkpoint: {e}")
    return False


# ============================================================================
# Workflow State Operations
# ============================================================================

def load_workflow_state(workflow_dir: str) -> dict:
    """Load current workflow state."""
    state_path = os.path.join(workflow_dir, 'workflow_state.json')
    return load_json(state_path)


def load_active_tasks(workflow_dir: str) -> List[dict]:
    """Load tasks that are in progress."""
    tasks_dir = os.path.join(workflow_dir, 'tasks')
    active_tasks = []

    if os.path.exists(tasks_dir):
        for filename in os.listdir(tasks_dir):
            if filename.endswith('.yml') or filename.endswith('.json'):
                task_path = os.path.join(tasks_dir, filename)
                task = load_json(task_path) if filename.endswith('.json') else {}
                if task.get('status') == 'in_progress':
                    active_tasks.append(task)

    return active_tasks


def get_pending_tasks(workflow_dir: str) -> List[dict]:
    """Get pending tasks in priority order."""
    tasks_dir = os.path.join(workflow_dir, 'tasks')
    pending = []

    if os.path.exists(tasks_dir):
        for filename in os.listdir(tasks_dir):
            if filename.endswith('.json'):
                task_path = os.path.join(tasks_dir, filename)
                task = load_json(task_path)
                if task.get('status') == 'pending':
                    pending.append(task)

    # Sort by layer, then by ID
    pending.sort(key=lambda t: (t.get('layer', 999), t.get('id', '')))
    return pending


# ============================================================================
# Context State Management
# ============================================================================

def capture_context_state(
    workflow_dir: str,
    context_percentage: float = 0.0,
    active_work: Optional[dict] = None,
    decisions: Optional[List[dict]] = None,
    blockers: Optional[List[dict]] = None
) -> dict:
    """Capture current context state for later resume."""

    workflow_state = load_workflow_state(workflow_dir)
    active_tasks = load_active_tasks(workflow_dir)
    pending_tasks = get_pending_tasks(workflow_dir)
    modified_files = get_git_status()

    # Determine active task
    active_task = active_tasks[0] if active_tasks else None

    # Build next actions from pending tasks
    next_actions = []
    for task in pending_tasks[:5]:  # Top 5 pending
        next_actions.append({
            'action': task.get('type', 'implement'),
            'target': task.get('title', task.get('id', 'unknown')),
            'priority': len(next_actions) + 1,
            'context_needed': [task.get('context', {}).get('context_snapshot_path', '')]
        })

    # Build context state
    state = {
        'session_id': f"compact_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        'captured_at': datetime.now().isoformat(),

        'context_usage': {
            'tokens_used': 0,  # Would need to be passed in
            'tokens_max': 0,
            'percentage': context_percentage,
            'threshold_triggered': THRESHOLDS['save']
        },

        'workflow_position': {
            'workflow_id': workflow_state.get('id', 'unknown'),
            'current_phase': workflow_state.get('current_phase', 'UNKNOWN'),
            'active_task_id': active_task.get('id') if active_task else None,
            'layer': active_task.get('layer', 1) if active_task else 1
        },

        'active_work': active_work or {
            'entity_id': active_task.get('entity_id', '') if active_task else '',
            'entity_type': active_task.get('type', '') if active_task else '',
            'action': 'implementing' if active_task else 'pending',
            'file_path': None,
            'progress_notes': ''
        },

        'next_actions': next_actions,
        'modified_files': modified_files,
        'decisions': decisions or [],
        'blockers': blockers or []
    }

    return state


def save_context_state(workflow_dir: str, state: dict):
    """Save context state to file."""
    state_path = os.path.join(workflow_dir, STATE_FILE)
    save_json(state_path, state)
    print(f"Saved context state to: {state_path}")

    # Also save modified files separately for quick access
    modified_path = os.path.join(workflow_dir, MODIFIED_FILES_FILE)
    save_json(modified_path, state.get('modified_files', []))


def generate_resume_prompt(state: dict) -> str:
    """Generate human-readable resume prompt."""

    lines = [
        "## Context Recovery - Resuming Previous Session",
        "",
        "### Session Info",
        f"- **Original Session**: {state.get('session_id', 'unknown')}",
        f"- **Captured At**: {state.get('captured_at', 'unknown')}",
        f"- **Context Usage**: {state.get('context_usage', {}).get('percentage', 0) * 100:.1f}%",
        "",
        "### Workflow Position",
        f"- **Phase**: {state.get('workflow_position', {}).get('current_phase', 'UNKNOWN')}",
        f"- **Active Task**: {state.get('workflow_position', {}).get('active_task_id', 'None')}",
        f"- **Layer**: {state.get('workflow_position', {}).get('layer', 1)}",
        "",
        "### What Was Being Worked On",
    ]

    active = state.get('active_work', {})
    lines.extend([
        f"- **Entity**: {active.get('entity_id', 'None')} ({active.get('entity_type', '')})",
        f"- **Action**: {active.get('action', 'unknown')}",
        f"- **File**: {active.get('file_path', 'None')}",
        f"- **Progress**: {active.get('progress_notes', 'No notes')}",
        "",
        "### Next Actions (Priority Order)",
    ])

    for action in state.get('next_actions', []):
        context = ', '.join(action.get('context_needed', [])) or 'None'
        lines.append(f"{action.get('priority', '?')}. **{action.get('action', '')}** {action.get('target', '')}")
        lines.append(f"   - Context needed: {context}")

    if state.get('modified_files'):
        lines.extend(["", "### Recent Changes"])
        for f in state.get('modified_files', []):
            lines.append(f"- `{f.get('path', '')}` - {f.get('action', '')}: {f.get('summary', '')}")

    if state.get('decisions'):
        lines.extend(["", "### Key Decisions Made"])
        for d in state.get('decisions', []):
            lines.append(f"- **{d.get('topic', '')}**: {d.get('decision', '')}")

    if state.get('blockers'):
        lines.extend(["", "### Current Blockers"])
        for b in state.get('blockers', []):
            lines.append(f"- {b.get('issue', '')} ({b.get('status', '')}): {b.get('notes', '')}")

    lines.extend([
        "",
        "---",
        "**Action Required**: Continue from the next action listed above.",
        "",
        "To load full context, read the following files:",
    ])

    # Add context files to read
    for action in state.get('next_actions', [])[:3]:
        for ctx in action.get('context_needed', []):
            if ctx:
                lines.append(f"- `{ctx}`")

    return '\n'.join(lines)


# ============================================================================
# Commands
# ============================================================================

def cmd_save(args):
    """Save current context state (pre-compact hook)."""
    workflow_dir = args.workflow_dir

    print("=" * 60)
    print("PRE-COMPACT: Saving Context State")
    print("=" * 60)

    # Capture state
    state = capture_context_state(
        workflow_dir,
        context_percentage=args.percentage / 100 if args.percentage else 0.80
    )

    # Save state
    save_context_state(workflow_dir, state)

    # Generate resume prompt
    resume_prompt = generate_resume_prompt(state)
    resume_path = os.path.join(workflow_dir, RESUME_PROMPT_FILE)
    save_text(resume_path, resume_prompt)
    print(f"Generated resume prompt: {resume_path}")

    # Create git checkpoint if requested
    if args.checkpoint:
        create_checkpoint("WIP: Pre-compaction checkpoint")

    print()
    print("=" * 60)
    print("State saved successfully!")
    print(f"  Session ID: {state['session_id']}")
    print(f"  Active Task: {state['workflow_position']['active_task_id']}")
    print(f"  Next Actions: {len(state['next_actions'])}")
    print("=" * 60)

    return 0


def cmd_resume(args):
    """Display resume prompt (post-compact hook)."""
    workflow_dir = args.workflow_dir

    # Load state
    state_path = os.path.join(workflow_dir, STATE_FILE)
    state = load_json(state_path)

    if not state:
        print("No saved context state found.")
        print(f"Expected: {state_path}")
        return 1

    # Generate and display resume prompt
    if args.json:
        print(json.dumps(state, indent=2))
    else:
        resume_prompt = generate_resume_prompt(state)
        print(resume_prompt)

    return 0


def cmd_status(args):
    """Show current context state status."""
    workflow_dir = args.workflow_dir

    print("=" * 60)
    print("CONTEXT STATE STATUS")
    print("=" * 60)

    # Check for saved state
    state_path = os.path.join(workflow_dir, STATE_FILE)
    if os.path.exists(state_path):
        state = load_json(state_path)
        print(f"\nSaved State Found:")
        print(f"  Session ID: {state.get('session_id', 'unknown')}")
        print(f"  Captured At: {state.get('captured_at', 'unknown')}")
        print(f"  Phase: {state.get('workflow_position', {}).get('current_phase', 'UNKNOWN')}")
        print(f"  Active Task: {state.get('workflow_position', {}).get('active_task_id', 'None')}")
    else:
        print(f"\nNo saved state at: {state_path}")

    # Check workflow state
    workflow_state = load_workflow_state(workflow_dir)
    if workflow_state:
        print(f"\nCurrent Workflow:")
        print(f"  ID: {workflow_state.get('id', 'unknown')}")
        print(f"  Phase: {workflow_state.get('current_phase', 'UNKNOWN')}")

    # Check git status
    modified = get_git_status()
    if modified:
        print(f"\nModified Files: {len(modified)}")
        for f in modified[:5]:
            print(f"  - {f['action']}: {f['path']}")
        if len(modified) > 5:
            print(f"  ... and {len(modified) - 5} more")

    return 0


def cmd_clear(args):
    """Clear saved context state."""
    workflow_dir = args.workflow_dir

    files_to_remove = [
        os.path.join(workflow_dir, STATE_FILE),
        os.path.join(workflow_dir, RESUME_PROMPT_FILE),
        os.path.join(workflow_dir, MODIFIED_FILES_FILE)
    ]

    removed = 0
    for filepath in files_to_remove:
        if os.path.exists(filepath):
            os.remove(filepath)
            print(f"Removed: {filepath}")
            removed += 1

    if removed == 0:
        print("No state files found to remove.")
    else:
        print(f"\nCleared {removed} state file(s).")

    return 0


# ============================================================================
# Main CLI
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Context Compaction Manager for Guardrail Orchestrator"
    )
    parser.add_argument(
        '--workflow-dir', '-w',
        default=DEFAULT_WORKFLOW_DIR,
        help='Workflow directory path'
    )

    subparsers = parser.add_subparsers(dest='command', help='Commands')

    # Save command
    save_parser = subparsers.add_parser('save', help='Save context state (pre-compact)')
    save_parser.add_argument('--percentage', '-p', type=float, help='Current context percentage (0-100)')
    save_parser.add_argument('--checkpoint', '-c', action='store_true', help='Create git checkpoint')

    # Resume command
    resume_parser = subparsers.add_parser('resume', help='Display resume prompt (post-compact)')
    resume_parser.add_argument('--json', '-j', action='store_true', help='Output as JSON')

    # Status command
    subparsers.add_parser('status', help='Show context state status')

    # Clear command
    subparsers.add_parser('clear', help='Clear saved context state')

    args = parser.parse_args()

    if args.command == 'save':
        return cmd_save(args)
    elif args.command == 'resume':
        return cmd_resume(args)
    elif args.command == 'status':
        return cmd_status(args)
    elif args.command == 'clear':
        return cmd_clear(args)
    else:
        parser.print_help()
        return 1


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