import argparse
import subprocess
import json
import ollama
from sttlib import load_whisper_model, record_until_enter, transcribe

# --- Configuration ---
OLLAMA_MODEL = "qwen2.5-coder:7b"
CONFIRM_COMMANDS = True  # Set to False to run commands instantly


# --- Terminal Tool ---

def run_terminal_command(command: str):
    """
    Executes a bash command in the Linux terminal.
    Used for file management, system info, and terminal tasks.
    """
    if CONFIRM_COMMANDS:
        print(f"\n{'='*40}")
        print(f"\u26a0\ufe0f  AI SUGGESTED: \033[1;32m{command}\033[0m")
        choice = input("   Confirm? [Y/n] or provide feedback: ").strip()

        if choice.lower() == 'n':
            return "USER_REJECTION: The user did not approve this command. Please suggest an alternative."
        elif choice and choice.lower() != 'y':
            return f"USER_FEEDBACK: The user rejected the command with this reason: '{choice}'. Please adjust."
        print(f"{'='*40}\n")

    # Safety Guardrail
    blacklist = ["rm -rf /", "mkfs", "dd if="]
    if any(forbidden in command for forbidden in blacklist):
        return "Error: Command blocked for security reasons."

    try:
        result = subprocess.run(command, shell=True,
                                capture_output=True, text=True, timeout=20)
        output = f"STDOUT: {result.stdout}\nSTDERR: {result.stderr}"
        return output if output.strip() else "Success (No output)."
    except Exception as e:
        return f"Execution Error: {str(e)}"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=OLLAMA_MODEL)
    parser.add_argument("--model-size", default="medium",
                        help="Whisper model size")
    args, _ = parser.parse_known_args()

    whisper_model = load_whisper_model(args.model_size)

    # Initial System Prompt
    messages = [{
        'role': 'system',
        'content': (
            'You are a Linux expert assistant. When asked for a system task, '
            'use the "run_terminal_command" tool. If the user rejects a command, '
            'analyze their feedback and suggest a corrected alternative.'
        )
    }]

    print(f"--- Assistant Active (Model: {args.model}) ---")

    while True:
        try:
            # 1. Voice Capture
            audio_data = record_until_enter()
            user_text = transcribe(whisper_model, audio_data.flatten())
            if not user_text:
                continue

            print(f"\nYOU: {user_text}")
            messages.append({'role': 'user', 'content': user_text})

            # 2. AI Interaction Loop (Supports up to 3 retries if rejected)
            for attempt in range(3):
                response = ollama.chat(
                    model=args.model,
                    messages=messages,
                    tools=[run_terminal_command],
                    options={'temperature': 0}
                )

                tool_calls = response.message.tool_calls

                # Fallback Repair: Catch raw JSON output
                if not tool_calls and '"run_terminal_command"' in response.message.content:
                    try:
                        c = response.message.content
                        start, end = c.find('{'), c.rfind('}') + 1
                        raw_json = json.loads(c[start:end])
                        tool_calls = [{'function': {
                            'name': 'run_terminal_command',
                            'arguments': raw_json.get('arguments', raw_json)
                        }}]
                    except:
                        pass

                # 3. Execution Logic
                if tool_calls:
                    call = tool_calls[0]
                    # Normalize arguments format
                    f_args = call.function.arguments if hasattr(
                        call, 'function') else call['function']['arguments']

                    result = run_terminal_command(f_args['command'])

                    # Update History
                    messages.append(response.message)
                    messages.append({'role': 'tool', 'content': result})

                    # If REJECTED, the loop continues and the AI sees the feedback
                    if "USER_REJECTION" in result or "USER_FEEDBACK" in result:
                        print("[RETHINKING] AI is adjusting the command...")
                        continue
                    else:
                        # Success: Let AI explain the result
                        final_res = ollama.chat(
                            model=args.model, messages=messages)
                        print(f"AI: {final_res.message.content}")
                        messages.append(final_res.message)
                        break
                else:
                    # Normal Chat
                    print(f"AI: {response.message.content}")
                    messages.append(response.message)
                    break

        except KeyboardInterrupt:
            print("\nExiting...")
            break
        except Exception as e:
            print(f"System Error: {e}")


if __name__ == "__main__":
    main()