# LLM agent player using HuggingFace's smolagent
from wiki_run_engine import WikiRunEnvironment
from rich.console import Console

console = Console()

try:
    from smolagent import Agent, AgentConfig
except ImportError:
    console.print("[red]smolagent package not found. Please install with 'uv pip install smolagent'[/red]")
    raise

class AgentPlayer:
    def __init__(self, wiki_data_path, model_name="HuggingFaceH4/zephyr-7b-beta"):
        """Initialize agent player"""
        self.env = WikiRunEnvironment(wiki_data_path)
        
        # Initialize LLM agent
        config = AgentConfig(
            model=model_name,
            output_parser="json"
        )
        self.agent = Agent(config)
        
    def play(self, start_article=None, target_article=None, max_steps=20):
        """Play a game of Wiki Run using the LLM agent"""
        # Reset environment
        state = self.env.reset(start_article, target_article)
        
        console.print("[bold]Agent Wiki Run[/bold]")
        console.print(f"Starting article: [cyan]{state['current_article']}[/cyan]")
        console.print(f"Target article: [red]{state['target_article']}[/red]")
        console.print()
        
        steps = 0
        while not state['is_complete'] and steps < max_steps:
            console.print(f"[bold]Step {steps + 1}:[/bold]")
            console.print(f"Current article: [cyan]{state['current_article']}[/cyan]")
            
            # Create prompt for agent
            prompt = self._create_agent_prompt(state)
            
            # Get agent's decision
            tool_result = self.agent.run(
                prompt,
                tools=[
                    {
                        "name": "choose_next_article",
                        "description": "Choose the next Wikipedia article to navigate to",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "article": {
                                    "type": "string",
                                    "description": "The title of the next article to navigate to"
                                },
                                "reasoning": {
                                    "type": "string",
                                    "description": "Explanation of why this article was chosen"
                                }
                            },
                            "required": ["article", "reasoning"]
                        }
                    }
                ]
            )
            
            # Extract agent's choice
            choice = tool_result.get("choose_next_article", {})
            next_article = choice.get("article", "")
            reasoning = choice.get("reasoning", "")
            
            console.print(f"[yellow]Agent reasoning: {reasoning}[/yellow]")
            console.print(f"Agent chooses: [green]{next_article}[/green]")
            
            # Verify the choice is valid
            if next_article in state['available_links']:
                state, message = self.env.step(next_article)
                if message:
                    console.print(f"[bold]{message}[/bold]")
            else:
                console.print("[red]Invalid choice! Agent selected an article that's not in the available links.[/red]")
                # Choose a random valid link as fallback
                import random
                next_article = random.choice(state['available_links'])
                console.print(f"[yellow]Falling back to random choice: {next_article}[/yellow]")
                state, _ = self.env.step(next_article)
                
            steps += 1
            console.print()
            
        # Game complete
        if state['is_complete']:
            console.print("[bold green]Success! Agent reached the target article![/bold green]")
        else:
            console.print("[bold red]Failed to reach target within step limit.[/bold red]")
            
        console.print(f"Steps taken: [bold]{state['steps_taken']}[/bold]")
        console.print(f"Path: [italic]{' → '.join(state['path_taken'])}[/italic]")
        
        return state
    
    def _create_agent_prompt(self, state):
        """Create prompt for the agent"""
        current = state['current_article']
        target = state['target_article']
        links = state['available_links']
        
        prompt = f"""You are playing the Wiki Run game. Your goal is to navigate from the current Wikipedia article to the target article using only the available links.

Current article: {current}
Target article: {target}

Available links (choose one):
{', '.join(links)}

Choose the link that you think will get you closest to the target article. Consider:
1. Direct connections to the target
2. Articles that might be in the same category as the target
3. General articles that might have many links to other topics

Use the choose_next_article tool to make your selection."""
        
        return prompt

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        console.print("[red]Please provide the path to Wikipedia data[/red]")
        console.print("Usage: python agent.py <wiki_data_path>")
        sys.exit(1)
        
    wiki_data_path = sys.argv[1]
    agent = AgentPlayer(wiki_data_path)
    agent.play()