import os
from pathlib import Path

from openhands.sdk import LLM, Agent, Conversation, Tool
from openhands.tools.file_editor import FileEditorTool
from openhands.tools.task_tracker import TaskTrackerTool
from openhands.tools.terminal import TerminalTool


# ============================================================================
# 基础示例 - 使用默认 system prompt
# ============================================================================
def basic_example():
    """使用默认的 system prompt"""
llm = LLM(
    model=os.getenv("LLM_MODEL", "vertex_ai/gemini-2.5-flash"),
    api_key=os.getenv("LLM_API_KEY"),
    base_url=os.getenv("LLM_BASE_URL", None),
)

agent = Agent(
    llm=llm,
    tools=[
        Tool(name=TerminalTool.name),
        Tool(name=FileEditorTool.name),
        Tool(name=TaskTrackerTool.name),
    ],
)

cwd = os.getcwd()
conversation = Conversation(agent=agent, workspace=cwd)
conversation.send_message("Write 3 facts about the current project into FACTS.txt.")
conversation.run()
print("All done!")


# ============================================================================
# 自定义示例 - 使用自定义的代码演化 prompt
# ============================================================================
def custom_evolution_example():
    """使用专门为代码演化任务定制的 system prompt"""
    llm = LLM(
        model=os.getenv("LLM_MODEL", "vertex_ai/gemini-2.5-flash"),
        api_key=os.getenv("LLM_API_KEY"),
        base_url=os.getenv("LLM_BASE_URL", None),
    )

    # 使用自定义的演化 prompt 模板
    custom_prompt_path = Path(__file__).parent / "custom_evolution_prompt.j2"
    
    agent = Agent(
        llm=llm,
        tools=[
            Tool(name=TerminalTool.name),
            Tool(name=FileEditorTool.name),
            Tool(name=TaskTrackerTool.name),
        ],
        # 指定自定义 prompt 文件
        system_prompt_filename=str(custom_prompt_path),
        # 传递额外的上下文变量到 prompt 模板
        system_prompt_kwargs={
            "custom_instructions": (
                "Focus on evolutionary improvements. "
                "Prioritize algorithmic efficiency and code clarity."
            ),
            "project_name": "ShinkaEvolve",
            "task_type": "Code Evolution",
            "project_context": "Evolutionary framework for optimizing scientific code",
            "task_description": "Analyze and improve the initial implementation",
        },
    )

    cwd = os.getcwd()
    conversation = Conversation(agent=agent, workspace=cwd)
    conversation.send_message(
        "Analyze the project structure and write 3 insights about "
        "potential code evolution strategies into FACTS.txt."
    )
    conversation.run()
    print("Evolution analysis complete!")


if __name__ == "__main__":
    # 运行基础示例
    # basic_example()
    
    # 或者运行自定义演化示例
    custom_evolution_example()