#基于 https://github.com/dave1010/hubcap 极简agent的python实现
import subprocess
import sys
import time

SYSTEM_PROMPT = """\
you are a coding agent. you can read and write files. Eg `cat helloworld.txt`, `echo "hello\nworld" > helloworld.txt`
output the next command required to progress your goal. output `DONE` when done."""

def chat(prompt: str, system: str | None = None) -> str:
    """Send a message to the LLM via the `llm` CLI tool."""
    options = ["-s", system] if system else ["-c"]
    print(f"\n\033[0;36m[PROMPT]\033[0m {prompt}")
    
    result = subprocess.run(
        ["llm", *options, prompt],
        capture_output=True,
        text=True
    )
    response = result.stdout.strip()
    print(f"\033[1;33m[RESPONSE]\033[0m {response}\n")
    return response


def run_command(cmd: str) -> tuple[int, str]:
    """Execute a shell command and return (return_code, output)."""
    result = subprocess.run(
        cmd,
        shell=True,
        capture_output=True,
        text=True
    )
    return result.returncode, result.stdout + result.stderr


def main():
    if len(sys.argv) < 2:
        print("Usage: python hubcap.py <goal>")
        sys.exit(1)

    goal = sys.argv[1]
    
    # Initial planning
    chat(f"GOAL: {goal}\n\nWHAT IS YOUR OVERALL PLAN?", SYSTEM_PROMPT)

    # Main loop: think → act → observe
    while True:
        response = chat(
            "SHELL COMMAND TO EXECUTE OR `DONE`. NO ADDITIONAL CONTEXT OR EXPLANATION:"
        ).strip()
        
        if response == "DONE":
            break

        print(f"\033[0;32m[EXECUTING]\033[0m {response}")
        time.sleep(3)  # Safety delay - allows user to Ctrl+C
        
        return_code, output = run_command(response)
        chat(
            f"COMMAND COMPLETED WITH RETURN CODE: {return_code}. OUTPUT:\n{output}\n\n"
            "WHAT ARE YOUR OBSERVATIONS?"
        )