#!/usr/bin/env python3
import sys, os, argparse, json, re, subprocess, signal, time, platform
from litellm import completion
from .config import AIOConfig, load_config, init_config_file, CONFIG_FILE

##############################################
# ANSI color codes
##############################################
class Colors:
    RESET = '\033[0m'
    BOLD = '\033[1m'
    DIM = '\033[2m'
    
    # Foreground colors
    RED = '\033[31m'
    GREEN = '\033[32m'
    YELLOW = '\033[33m'
    BLUE = '\033[34m'
    MAGENTA = '\033[35m'
    CYAN = '\033[36m'
    WHITE = '\033[37m'
    
    # Bright colors
    BRIGHT_RED = '\033[91m'
    BRIGHT_GREEN = '\033[92m'
    BRIGHT_YELLOW = '\033[93m'
    BRIGHT_BLUE = '\033[94m'
    BRIGHT_MAGENTA = '\033[95m'
    BRIGHT_CYAN = '\033[96m'

##############################################
# Output helpers
##############################################
def err(*args, color='', **kwargs):
    """Write to stderr with optional color."""
    if color:
        print(f"{color}{args[0]}{Colors.RESET}", *args[1:], file=sys.stderr, **kwargs)
    else:
        print(*args, file=sys.stderr, **kwargs)

def out(*args, color='', **kwargs):
    """Write to stdout with optional color."""
    if color:
        print(f"{color}{args[0]}{Colors.RESET}", *args[1:], **kwargs)
    else:
        print(*args, **kwargs)

##############################################
# Default config values
##############################################

def get_system_prompt(max_lines, max_chars, interactive=True, system_info=""):
    """Generate system prompt with truncation limits and mode-specific guidance."""
    
    # Base prompt (both modes)
    base = (
        f"You are an AI assistant operating in a terminal environment.\n"
        f"System: {system_info}\n"
        f"Your role is to help users accomplish tasks by providing shell commands when appropriate.\n"
        f"\nContext Awareness:"
        f"\n- Input buffers/files are truncated to {max_lines} lines / {max_chars} chars before being sent to you"
        f"\n- Command outputs are truncated to {max_lines} lines / {max_chars} chars when sent back to you"
        f"\n- When buffers are truncated, note this and work with available information"
        f"\n\nGuidelines:"
        f"\n- For file/data operations: provide the necessary bash commands rather than explaining what to do"
        f"\n- For general knowledge questions or tasks requiring explanation: answer directly without commands"
        f"\n- For intermediate/temporary files: use /tmp directory unless user specifies otherwise"
        f"\n- Always use --help to check command usage if unsure or ambiguous (in particular for aio/ai)"
        f"\n- Consider command safety - mention if a command modifies/deletes files"
        f"\n- Be concise and practical in your responses"
    )
    
    if interactive:
        # Interactive mode: user can execute commands and continue conversation
        interactive_addition = (
            f"\n- User manually executes commands you suggest - they are not auto-executed"
            f"\n- After execution, command output is sent back to you for analysis"
            f"\n\nCommand Generation Rules:"
            f"\n- Provide exactly ONE command block unless multiple sequential steps are truly necessary"
            f"\n- Maximum 3 commands total (only for complex multi-step tasks)"
            f"\n- Each command must be in a separate ```bash code block"
            f"\n- Start each block with a # comment explaining what the command does"
            f"\n- Users can execute commands by typing the number (1, 2, or 3)"
            f"\n\nRequired format:"
            f"\n```bash"
            f"\n# Clear explanation of what this command does"
            f"\nactual command here"
            f"\n```"
            f"\n\nThis format is parsed automatically - follow it exactly."
        )
        return base + interactive_addition
    else:
        # Non-interactive mode: one-shot interaction
        exec_addition = (
            f"\n\nExecution Mode:"
            f"\n- This is a one-shot interaction - the user cannot respond or execute commands after your response"
            f"\n- When appropriate, provide commands in ```bash code blocks with # comment explanations"
            f"\n- Generate commands for: data processing, file operations, log analysis, text manipulation"
            f"\n- Provide direct answers for: general knowledge, explanations, summaries, questions requiring reasoning"
            f"\n- Common use cases: analyzing logs, parsing JSON/HTML/CSV files, processing large text files"
        )
        return base + exec_addition

##############################################
# CLI argument parsing
##############################################

def parse_bash_commands(text):
    """Extract bash code blocks from markdown text with explanations.
    
    Returns:
        list of tuples: [(explanation, command), ...]
    """
    # Match both ```bash and ``` (without language specifier)
    pattern = r'```(?:bash)?\s*\n(.*?)\n```'
    matches = re.findall(pattern, text, re.DOTALL)
    
    results = []
    for match in matches[:5]:  # Max 5 commands
        lines = match.strip().split('\n')
        explanation = ""
        command_lines = []
        
        for line in lines:
            stripped = line.strip()
            if stripped.startswith('#'):
                # Extract explanation from comment
                if not explanation:
                    explanation = stripped[1:].strip()
            else:
                command_lines.append(line)
        
        command = '\n'.join(command_lines).strip()
        if command:
            if not explanation:
                explanation = "Execute command"  # Default if no explanation provided
            results.append((explanation, command))
    
    return results

def truncate_buffer(raw_buffer, max_lines, max_chars, filename=None):
    """
    Truncate buffer to specified limits and return truncated buffer with info.
    
    Returns:
        tuple: (truncated_buffer, info_string)
    """
    if not raw_buffer.strip():
        return None, ""
    
    lines = raw_buffer.splitlines()
    original_lines = len(lines)
    original_chars = len(raw_buffer)
    
    truncated_lines = lines[:max_lines]
    truncated_text = "\n".join(truncated_lines)
    truncated_text = truncated_text[:max_chars]
    
    truncated_lines_count = len(truncated_text.splitlines())
    truncated_chars_count = len(truncated_text)
    
    # Build info string
    info = ""
    if filename:
        info = f"File: {filename}, "
    
    if original_lines != truncated_lines_count or original_chars != truncated_chars_count:
        info += f"truncated from {original_lines} lines / {original_chars} chars to {truncated_lines_count} lines / {truncated_chars_count} chars"
    else:
        info += f"{truncated_lines_count} lines / {truncated_chars_count} chars"
    
    return truncated_text, info

def execute_command(command, max_lines, max_chars, timeout=60):
    """
    Execute a shell command and return the output.
    
    Args:
        command: The shell command to execute
        max_lines: Maximum lines for truncation context
        max_chars: Maximum chars for truncation context
        timeout: Command timeout in seconds
    
    Returns:
        tuple: (success: bool, output: str, truncated_output: str, return_code: int)
    """
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout
        )
        
        output = result.stdout
        if result.stderr:
            output += "\nSTDERR:\n" + result.stderr
        
        exit_color = Colors.GREEN if result.returncode == 0 else Colors.RED
        err(f"[Exit code: {result.returncode}]", color=exit_color)
        out(output)
        
        # Truncate output for LLM context
        truncated_output, output_info = truncate_buffer(
            output, 
            max_lines=max_lines,
            max_chars=max_chars
        )
        
        if output_info and "truncated" in output_info:
            err(f"[Output truncated for LLM context: {output_info}]", color=Colors.DIM)
        
        return True, output, truncated_output or output, result.returncode
        
    except subprocess.TimeoutExpired:
        err(f"[Command timed out after {timeout} seconds]", color=Colors.RED)
        return False, "", f"Error: Command timed out after {timeout} seconds", -1
    except Exception as e:
        err(f"[Execution error: {e}]", color=Colors.RED)
        return False, "", f"Error: {e}", -1

def main_wrapper():
    # Check for special commands (init) before normal arg parsing
    if len(sys.argv) > 1 and sys.argv[1] == "init":
        # Handle init command
        force = "--force" in sys.argv or "-f" in sys.argv
        
        if init_config_file(force=force):
            out(f"Configuration file created: {CONFIG_FILE}", color=Colors.GREEN)
            out(f"\nEdit this file to customize your defaults.", color=Colors.DIM)
            default_config = AIOConfig()
            out(f"Current model: {default_config.model}", color=Colors.DIM)
        else:
            err(f"Configuration file already exists: {CONFIG_FILE}", color=Colors.YELLOW)
            err(f"Use 'ai init --force' to overwrite.", color=Colors.DIM)
        sys.exit(0)

    parser = argparse.ArgumentParser(
        description="AI terminal helper using LiteLLM. "
                    "Reads stdin optionally, summarizes or processes content.\n\n"
                    "Use 'ai init' to create a config file with your defaults.",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument("-m", "--model", type=str,
                        help="Model name (default: from config file, overriden by env AIO_MODEL_NAME)")

    parser.add_argument("-n", "--lines", type=int,
                        help="Max input lines (default: from config file)")

    parser.add_argument("-c", "--chars", type=int,
                        help="Max input characters (default: from config file)")

    parser.add_argument("-r", "--reasoning-effort", type=str, # allow all / choices=["low", "medium", "high"],
                        help="Reasoning effort for thinking models (default: from config file)")

    parser.add_argument("-f", "--file", type=str,
                        help="Read input from file instead of stdin")

    parser.add_argument("-d", "--dry-run", action="store_true",
                        help="Show what would be sent to litellm without making the API call")

    parser.add_argument("-e", "--exec", action="store_true",
                        help="Non-interactive mode: execute once and exit (for pipes)")

    parser.add_argument("instruction", nargs="*", help="Instruction text (if none & stdin/file present, summarize buffer)")

    args = parser.parse_args()

    ##############################################
    # Load configuration with tiered priority
    # Priority: CLI args > Env vars > Config file > Defaults
    ##############################################
    # 1. Load from config file or defaults
    config = load_config()
    
    # 2. Apply overrides: CLI > ENV > Config
    model = args.model or os.getenv("AIO_MODEL_NAME") or config.model
    
    max_lines = args.lines
    if max_lines is None:
        env_lines = os.getenv("AIO_MAX_LINES")
        max_lines = int(env_lines) if env_lines else config.max_lines
    
    max_chars = args.chars
    if max_chars is None:
        env_chars = os.getenv("AIO_MAX_CHARS")
        max_chars = int(env_chars) if env_chars else config.max_chars
    
    reasoning_effort = args.reasoning_effort
    if reasoning_effort is None:
        env_reasoning = os.getenv("AIO_REASONING_EFFORT")
        reasoning_effort = env_reasoning if env_reasoning else config.reasoning_effort

    ##############################################
    # Read from file or stdin if present
    ##############################################
    raw_buffer = ""
    stdin_was_piped = False
    if args.file:
        try:
            with open(args.file, 'r') as f:
                raw_buffer = f.read()
        except Exception as e:
            err(f"[ai_cmd] Error reading file: {e}", color=Colors.RED)
            sys.exit(1)
    elif not sys.stdin.isatty():
        raw_buffer = sys.stdin.read()
        stdin_was_piped = True
    
    # Auto-enable exec mode if stdin was piped (stdin consumed, can't be interactive)
    if stdin_was_piped:
        args.exec = True

    ##############################################
    # Prepare buffer if present
    ##############################################
    buffer, buffer_info = truncate_buffer(raw_buffer, max_lines, max_chars, args.file)
    
    # Display buffer info to stderr if we have a buffer
    if buffer_info:
        err(f"[ai_cmd] {buffer_info}", color=Colors.DIM)

    ##############################################
    # Determine instruction behavior
    ##############################################
    if buffer and not args.instruction:
        # Has buffer, no instruction: summarize
        instruction = "Provide a quick summary of the provided buffer."
    elif args.instruction:
        # Has instruction
        instruction = " ".join(args.instruction)
    elif buffer:
        # Has buffer with instruction (already handled above)
        pass
    else:
        # No buffer, no instruction: get first message from user
        instruction = None

    ##############################################
    # Build initial message (if we have an instruction)
    ##############################################
    if instruction:
        user_message = instruction.strip()
        if buffer:
            # Add buffer to user message with info
            user_message += "\n\nBUFFER"
            if args.file:
                user_message += f" (from file: {args.file})"
            user_message += ":\n" + buffer.strip()
            user_message += f"\n\n({buffer_info})"
    else:
        # No initial message, will prompt user in interactive mode
        user_message = None


    ##############################################
    # Call OpenRouter using litellm with streaming
    ##############################################
    
    # Detect system information
    system_os = platform.system()
    system_arch = platform.machine()
    system_shell = os.path.basename(os.environ.get('SHELL', 'cmd' if system_os == 'Windows' else 'sh'))
    system_cwd = os.getcwd()
    system_info = f"OS:{system_os} Arch:{system_arch} Shell:{system_shell} CWD:{system_cwd}"
    
    # Display system info and config to user
    err(f"[ai_cmd] System: {system_info}", color=Colors.DIM)
    err(f"[ai_cmd] Model: {model}", color=Colors.DIM)
    
    # Initialize conversation history with appropriate system prompt
    system_prompt = get_system_prompt(
        max_lines, max_chars,
        interactive=not args.exec,
        system_info=system_info
    )
    
    messages = [{"role": "system", "content": system_prompt}]
    
    # Track available commands for execution in interactive mode
    available_commands = {}
    
    # Add initial user message if we have one
    if user_message:
        messages.append({"role": "user", "content": user_message})
    
    # If dry-run, just print the parameters and exit
    if args.dry_run:
        dry_run_output = {
            "model": model,
            "messages": messages,
            "stream": True
        }
        print(json.dumps(dry_run_output, indent=2))
        sys.exit(0)
    
    try:
        # Control-C interrupt tracking
        interrupt_count = 0
        last_interrupt_time = 0
        
        def signal_handler(sig, frame):
            nonlocal interrupt_count, last_interrupt_time
            current_time = time.time()
            
            # Reset interrupt count if it's been a while since last interrupt
            if current_time - last_interrupt_time > 2:
                interrupt_count = 0
            
            interrupt_count += 1
            last_interrupt_time = current_time
            
            if interrupt_count >= 2:
                err("\n[Exiting...]", color=Colors.YELLOW)
                sys.exit(0)
            else:
                err("\n[Interrupted - press Ctrl+C again to exit, or continue typing]", color=Colors.YELLOW)
                raise KeyboardInterrupt()
        
        signal.signal(signal.SIGINT, signal_handler)
        
        # Show welcome message in interactive mode (only if no initial user message)
        if not args.exec and not user_message:
            err("AI Terminal Assistant", color=Colors.BRIGHT_CYAN)
            # Display compact config
            config_parts = [
                f"model:{model}",
                f"lines:{max_lines}",
                f"chars:{max_chars}"
            ]
            if reasoning_effort:
                config_parts.append(f"reasoning:{reasoning_effort}")
            err(f"[{' | '.join(config_parts)}]", color=Colors.DIM)
            err("Commands:", color=Colors.BRIGHT_GREEN)
            err("  • Type 1-3 to execute AI-suggested commands", color=Colors.DIM)
            err("  • Type !<command> to execute shell commands directly (e.g., !ls -la)", color=Colors.DIM)
            err("  • Type /reset to clear conversation history", color=Colors.DIM)
            err("  • Press Ctrl+C twice to exit", color=Colors.DIM)
            err("Ask anything or request commands:", color=Colors.BRIGHT_GREEN)
        
        # Interactive conversation loop
        while True:
            # Get user input if we don't have a pending user message
            last_message_role = messages[-1]["role"] if len(messages) > 1 else "system"
            if last_message_role != "user":
                try:
                    # Show available commands if any
                    if available_commands and not args.exec:
                        err(f"[Commands available: {', '.join(available_commands.keys())}]", color=Colors.YELLOW)
                    
                    prompt_color = Colors.BRIGHT_BLUE if last_message_role == "assistant" else Colors.BLUE
                    err("> " if last_message_role == "assistant" else "> ", color=prompt_color, end='')
                    user_input = input()
                    
                    # Reset interrupt count on successful input
                    interrupt_count = 0
                    
                    if not user_input.strip():
                        continue
                    
                    # Check for /reset command
                    if user_input.strip() == "/reset":
                        messages = [{"role": "system", "content": system_prompt}]
                        available_commands = {}
                        err("[Chat reset]", color=Colors.GREEN)
                        continue
                    
                    # Determine what commands to execute
                    commands_to_run = []  # List of (label, command) tuples
                    
                    # Check for direct command execution with ! prefix
                    if user_input.strip().startswith("!"):
                        command = user_input.strip()[1:]  # Remove ! prefix
                        if not command:
                            err("[Error: No command specified after !]", color=Colors.RED)
                            continue
                        commands_to_run.append(("direct command", command))
                    
                    # Check if user wants to execute numbered command(s) (e.g., "1" or "12" or "123")
                    elif not args.exec and user_input.strip().isdigit() and available_commands and all(c in available_commands for c in user_input.strip()):
                        for cmd_num in user_input.strip():
                            explanation, command = available_commands[cmd_num]
                            commands_to_run.append((f"command {cmd_num}: {explanation}", command))
                    
                    # Execute commands if any were identified
                    if commands_to_run:
                        if len(commands_to_run) == 1:
                            err(f"[Executing {commands_to_run[0][0]}]", color=Colors.BRIGHT_CYAN)
                        else:
                            err(f"[Executing {len(commands_to_run)} commands sequentially]", color=Colors.BRIGHT_CYAN)
                        
                        exec_records = []
                        for label, command in commands_to_run:
                            if len(commands_to_run) > 1:
                                err(f"\n[{label}]", color=Colors.BRIGHT_CYAN)
                            err(command, color=Colors.CYAN)
                            
                            # Execute the command
                            success, output, truncated_output, return_code = execute_command(
                                command, max_lines, max_chars
                            )
                            
                            # Build execution record
                            record = f"{label}:\n{command}\n\nOutput:\n{truncated_output}"
                            if return_code != 0:
                                record += f"\n\nExit code: {return_code}"
                            exec_records.append(record)
                        
                        # Add execution record(s) to conversation history
                        if len(exec_records) == 1:
                            exec_record = f"User executed {exec_records[0]}"
                        else:
                            exec_record = f"User executed commands sequentially:\n\n" + "\n\n".join(exec_records)
                        
                        messages.append({"role": "user", "content": exec_record})
                        continue  # Go to next iteration to get AI response
                    
                    messages.append({"role": "user", "content": user_input})
                except KeyboardInterrupt:
                    # Handled by signal handler
                    continue
                except EOFError:
                    err()
                    break
            
            # Prepare call parameters
            call_params = {
                "model": model,
                "messages": messages,
                "stream": True
            }
            
            # Add reasoning_effort if specified (for reasoning models)
            if reasoning_effort:
                call_params["reasoning_effort"] = reasoning_effort
            else:  # set think to false (for ollama, see https://ollama.com/blog/thinking)
                call_params["think"] = False
            
            # Get assistant response
            response = completion(**call_params)
            assistant_message = ""
            reasoning_content = ""
            
            try:
                for chunk in response:
                    # Print content as it streams
                    if hasattr(chunk, 'choices') and len(chunk.choices) > 0:  # type: ignore
                        delta = chunk.choices[0].delta  # type: ignore
                        
                        # Handle reasoning content (thinking) separately
                        if hasattr(delta, 'reasoning_content') and delta.reasoning_content:
                            print(f"{Colors.DIM}{delta.reasoning_content}{Colors.RESET}", end='', flush=True)
                            reasoning_content += delta.reasoning_content
                        
                        # Handle regular content
                        if hasattr(delta, 'content') and delta.content:
                            print(delta.content, end='', flush=True)
                            assistant_message += delta.content
            except KeyboardInterrupt:
                err("\n[Response interrupted]", color=Colors.YELLOW)
                # Continue with whatever partial message we got
                if not assistant_message:
                    continue  # Skip if we got nothing
            
            print()  # Final newline
            
            # Add assistant response to conversation history
            messages.append({"role": "assistant", "content": assistant_message})
            
            # Parse bash commands from response (interactive mode only)
            if not args.exec:
                commands = parse_bash_commands(assistant_message)
                available_commands = {str(i+1): cmd_tuple for i, cmd_tuple in enumerate(commands)}
                if available_commands:
                    err("\n[Available commands:]", color=Colors.BRIGHT_GREEN)
                    for num, (explanation, cmd) in available_commands.items():
                        # Show explanation and command preview
                        cmd_preview = cmd.split('\n')[0][:60]
                        if len(cmd.split('\n')[0]) > 60 or '\n' in cmd:
                            cmd_preview += "..."
                        err(f"  {num}: {explanation}", color=Colors.GREEN)
                        err(f"     {cmd_preview}", color=Colors.DIM)
            
            # Exit if in exec mode
            if args.exec:
                break
            
    except Exception as e:
        err(f"[ai_cmd] Error: {e}", color=Colors.RED)
        raise e

if __name__ == "__main__":
    main_wrapper()