#!/usr/bin/env python3
"""Enhanced autorun plugin - supports arbitrary command wrapping and enhanced prompt injection"""

import json
import sys
import os
from pathlib import Path

# Add the src directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

try:
    from autorun.command_discovery import (
        discover_existing_commands,
        load_command_content,
        parse_command_args,
        validate_command_exists,
        get_command_metadata,
        find_command_by_pattern
    )
    COMMAND_DISCOVERY_AVAILABLE = True
except ImportError:
    # Fallback if command discovery not available
    COMMAND_DISCOVERY_AVAILABLE = False

# Import centralized configuration (DRY principle - single source of truth)
try:
    from autorun.config import CONFIG as BASE_CONFIG
    # Extend with command-specific mappings
    CONFIG = dict(BASE_CONFIG)
    CONFIG["command_mappings"] = dict(BASE_CONFIG["command_mappings"])
    CONFIG["command_mappings"].update({
        "/autorun-discover": "command_discovery",
        "/autorun-search": "command_search"
    })
    # Map legacy key names for backward compatibility
    CONFIG["completion_marker"] = BASE_CONFIG["stage1_confirmation"]
    CONFIG["emergency_stop_phrase"] = BASE_CONFIG["emergency_stop"]
except ImportError:
    # Fallback configuration if config.py cannot be imported
    CONFIG = {
        "command_mappings": {
            "/autorun": "activate",
            "/autoproc": "activate",
            "/autostop": "stop",
            "/estop": "emergency_stop",
            "/afs": "SEARCH",
            "/afa": "ALLOW",
            "/afj": "JUSTIFY",
            "/afst": "STATUS",
            "/autorun-discover": "command_discovery",
            "/autorun-search": "command_search"
        },
        "policies": {
            "ALLOW": ("allow-all", "ALLOW ALL: Full permission to create/modify files."),
            "JUSTIFY": ("justify-create", "JUSTIFIED: Search existing first. Include <AUTOFILE_JUSTIFICATION>reason</AUTOFILE_JUSTIFICATION> for new files."),
            "SEARCH": ("strict-search", "STRICT SEARCH: ONLY modify existing files. Use {glob} and {grep}. NO new files.")
        },
        "completion_marker": "AUTORUN_ALL_TASKS_COMPLETED_AND_VERIFIED_SUCCESSFULLY",
        "emergency_stop_phrase": "AUTORUN_STATE_PRESERVATION_EMERGENCY_STOP",
        "injection_template": """Your primary objective is to continue **UNINTERRUPTED, FULLY AUTONOMOUS, NONINTERACTIVE, PATIENT, AND SAFE EXECUTION** of your current tasks and goals.

1.  **MANDATORY PROCESS TO CONTINUE EXECUTION:** Cautiously and deliberately continue working **carefully, patiently, concretely, and safely**, non-stop, autonomously, and non-interactively per your instructions.
    * **Permissions Granted:** You have full permission to use all tools without requiring further permission prompts.

FILE CREATION POLICY: {policy_instructions}

2.  **SYSTEM STOP SIGNAL RULE:** The exact strings **{emergency_stop_phrase}** and **{completion_marker}** are **SYSTEM STOP SIGNALS**. You MUST **NEVER** output these strings unless corresponding stop condition has been met.

3.  **Safety Protocol:** Execute full Preservation Sequence **ONLY IF** current task involves high-risk irreversible actions.

4.  **FINAL OUTPUT ON SUCCESS:** When all tasks are 100% complete, output: **{completion_marker}**

Original task: {original_task}"""
    }

# Simple state using dict for now - avoid shelve complexity
_simple_state = {}

def handle_search():
    """Handle SEARCH command"""
    policy_name, policy_desc = CONFIG["policies"]["SEARCH"]
    return f"AutoFile policy: {policy_name} - {policy_desc}"

def handle_allow():
    """Handle ALLOW command"""
    policy_name, policy_desc = CONFIG["policies"]["ALLOW"]
    return f"AutoFile policy: {policy_name} - {policy_desc}"

def handle_justify():
    """Handle JUSTIFY command"""
    policy_name, policy_desc = CONFIG["policies"]["JUSTIFY"]
    return f"AutoFile policy: {policy_name} - {policy_desc}"

def handle_status():
    """Handle STATUS command"""
    current_policy = _simple_state.get("file_policy", "ALLOW")
    policy_name, policy_desc = CONFIG["policies"][current_policy]
    return f"Current policy: {policy_name}"

def handle_stop():
    """Handle STOP command"""
    return "Autorun stopped"

def handle_emergency_stop():
    """Handle EMERGENCY_STOP command"""
    return "Emergency stop activated"

def handle_activate(prompt):
    """Handle AUTORUN activation"""
    # Store original prompt
    _simple_state["activation_prompt"] = prompt
    _simple_state["file_policy"] = _simple_state.get("file_policy", "ALLOW")

    policy = _simple_state["file_policy"]
    policy_instructions = CONFIG["policies"][policy][1]

    # Use the shared injection template
    return CONFIG["injection_template"].format(
        emergency_stop_phrase=CONFIG["emergency_stop_phrase"],
        completion_marker=CONFIG["completion_marker"],
        policy_instructions=policy_instructions,
        original_task=prompt
    )

def handle_wrapped_command(prompt, original_command, args=""):
    """Handle arbitrary command wrapped with autorun functionality"""
    if not COMMAND_DISCOVERY_AVAILABLE:
        return f"Command discovery not available for: {original_command}"

    # Get command metadata
    cmd_info = get_command_metadata(original_command)
    if not cmd_info:
        return f"Command not found: {original_command}"

    # Load original command content
    original_content = load_command_content(cmd_info)
    if not original_content:
        return f"Error loading command content for: {original_command}"

    # Store original prompt and command for later use
    _simple_state["activation_prompt"] = prompt
    _simple_state["wrapped_command"] = original_command
    _simple_state["wrapped_args"] = args
    _simple_state["original_command_content"] = original_content
    _simple_state["file_policy"] = _simple_state.get("file_policy", "ALLOW")

    policy = _simple_state["file_policy"]
    policy_instructions = CONFIG["policies"][policy][1]

    # Create combined task description that includes both the wrapped command and original request
    combined_task = f"""AUTORUN WRAPPED COMMAND: {original_command}

=== ORIGINAL COMMAND CONTENT ===
{original_content}
=== END ORIGINAL COMMAND ===

Command arguments: {args if args else "None"}

Original wrapper request: {prompt}

AUTORUN INSTRUCTIONS:
Execute the above command with full autorun functionality while preserving all original command behavior.
- Execute the original command as specified above
- Handle any arguments provided
- Ensure all functionality of the original command is preserved
- Apply autorun methodology to extend autonomous execution
- Do not simply echo or repeat the command - actually execute it"""

    # Use the shared injection template with the combined task
    injection = CONFIG["injection_template"].format(
        emergency_stop_phrase=CONFIG["emergency_stop_phrase"],
        completion_marker=CONFIG["completion_marker"],
        policy_instructions=policy_instructions,
        original_task=combined_task
    )

    return injection

def handle_command_discovery():
    """Handle command discovery functionality"""
    if not COMMAND_DISCOVERY_AVAILABLE:
        return "Command discovery not available"

    commands = discover_existing_commands()
    stats = {
        "total": len(commands),
        "by_type": {},
        "by_source": {}
    }

    for cmd_info in commands.values():
        cmd_type = cmd_info.get("type", "unknown")
        source = cmd_info.get("source", "unknown")

        stats["by_type"][cmd_type] = stats["by_type"].get(cmd_type, 0) + 1
        stats["by_source"][source] = stats["by_source"].get(source, 0) + 1

    result = f"Discovered {stats['total']} commands:\n"
    result += f"By type: {dict(stats['by_type'])}\n"
    result += f"By source: {dict(stats['by_source'])}"

    return result

def handle_command_search(query):
    """Handle command search functionality"""
    if not COMMAND_DISCOVERY_AVAILABLE:
        return "Command discovery not available"

    matches = find_command_by_pattern(query, limit=10)
    if not matches:
        return f"No commands found matching: {query}"

    result = f"Found {len(matches)} commands matching '{query}':\n"
    for match in matches:
        cmd_name = match["command"]
        source = match.get("source", "unknown")
        score = match["score"]
        result += f"  {cmd_name} (source: {source}, score: {score})\n"

    return result

# Command handlers
COMMAND_HANDLERS = {
    "SEARCH": handle_search,
    "ALLOW": handle_allow,
    "JUSTIFY": handle_justify,
    "STATUS": handle_status,
    "stop": handle_stop,
    "emergency_stop": handle_emergency_stop,
    "activate": handle_activate,
    "wrapped_command": handle_wrapped_command,
    "command_discovery": handle_command_discovery,
    "command_search": handle_command_search
}

def main():
    """Standalone plugin entry point"""
    try:
        # Read input from stdin
        input_data = sys.stdin.read()

        if not input_data.strip():
            result = {
                "continue": True,
                "response": "",
                "error": "No input provided"
            }
            print(json.dumps(result, sort_keys=True))
            sys.stdout.flush()
            return

        # Parse JSON input
        try:
            payload = json.loads(input_data)
        except json.JSONDecodeError as e:
            result = {
                "continue": True,
                "response": "",
                "error": f"Invalid JSON: {e}"
            }
            print(json.dumps(result, sort_keys=True))
            sys.stdout.flush()
            return

        # Extract input data
        prompt = payload.get('prompt', '').strip()
        session_id = payload.get('session_id', 'default')

        # Enhanced command detection with arbitrary command support
        command = None
        wrapped_command = None
        command_args = ""

        # Try exact match first for known autorun commands
        command = CONFIG["command_mappings"].get(prompt, None)

        if not command:
            # Check for commands that support arguments (autorun/autoproc and prefixed versions)
            for cmd_key, cmd_value in CONFIG["command_mappings"].items():
                if (cmd_key.startswith('/autorun') or cmd_key.startswith('/autoproc')) and prompt.startswith(cmd_key):
                    command = cmd_value
                    break

        # Handle arbitrary command wrapping: /autorun:yourcommand
        if not command and COMMAND_DISCOVERY_AVAILABLE and prompt.startswith('/autorun:'):
            # Extract command name after /autorun:
            wrapped_command = prompt[9:].strip()  # Remove '/autorun:' prefix

            # Handle potential arguments in the wrapped command
            if ' ' in wrapped_command:
                parts = wrapped_command.split(' ', 1)
                wrapped_command = parts[0]
                command_args = parts[1] if len(parts) > 1 else ""

            # Add leading slash back for command lookup
            wrapped_command = f"/{wrapped_command}"

            # Validate that the wrapped command exists
            if validate_command_exists(wrapped_command):
                command = "wrapped_command"
            else:
                wrapped_command = None

        # Handle space-separated autorun commands: /autorun /yourcommand
        if not command and prompt.startswith('/autorun ') and len(prompt) > 10:
            space_separated_cmd = prompt[10:].strip()  # Remove '/autorun ' prefix

            # Handle potential arguments
            if ' ' in space_separated_cmd:
                parts = space_separated_cmd.split(' ', 1)
                space_separated_cmd = parts[0]
                command_args = parts[1] if len(parts) > 1 else ""

            # Validate that the command exists
            if validate_command_exists(space_separated_cmd):
                wrapped_command = space_separated_cmd
                command = "wrapped_command"

        # Handle autorun search command: /autorun-search query
        if not command and prompt.startswith('/autorun-search ') and len(prompt) > 17:
            search_query = prompt[17:].strip()  # Remove '/autorun-search ' prefix
            if search_query:
                response = handle_command_search(search_query)
                result = {
                    "continue": False,  # Command handled locally
                    "response": response
                }
                print(json.dumps(result, sort_keys=True))
                sys.stdout.flush()
                return

        if command and command in COMMAND_HANDLERS:
            # Handle command locally
            try:
                if command == "activate":
                    response = COMMAND_HANDLERS[command](prompt)
                elif command == "wrapped_command" and wrapped_command:
                    response = COMMAND_HANDLERS[command](prompt, wrapped_command, command_args)
                else:
                    response = COMMAND_HANDLERS[command]()

                # Update state for policy commands
                if command in ["SEARCH", "ALLOW", "JUSTIFY"]:
                    _simple_state["file_policy"] = command

                result = {
                    "continue": False,  # Command handled locally
                    "response": response
                }
            except Exception as e:
                result = {
                    "continue": True,
                    "response": "",
                    "error": f"Command execution failed: {e}"
                }
        else:
            # Let AI handle non-commands
            result = {
                "continue": True,  # Pass to AI
                "response": ""
            }

        # Return JSON response
        print(json.dumps(result, sort_keys=True))
        sys.stdout.flush()

    except Exception as e:
        result = {
            "continue": True,
            "response": "",
            "error": f"Unexpected error: {e}"
        }
        print(json.dumps(result, sort_keys=True))
        sys.stdout.flush()

if __name__ == "__main__":
    main()
