import litellm
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.table import Table
from rich.box import SIMPLE
from wiki_run_engine import WikiRunEnvironment
from langsmith import traceable
import os
from openai import OpenAI
from langsmith.wrappers import wrap_openai

openai_client = wrap_openai(OpenAI())

class WikiRunAgent:
    def __init__(self, wiki_data_path, model="gemini/gemini-2.5-pro-exp-03-25"):
        self.env = WikiRunEnvironment(wiki_data_path)
        self.model = model
        self.console = Console()
        
    @traceable(name="WikiRun Game")
    def run_game(self, start_article=None, target_article=None):
        """Play the WikiRun game with LLM agent"""
        state = self.env.reset(start_article, target_article)
        
        self.console.print(Panel(f"[bold cyan]Starting WikiRun![/bold cyan]"))
        self.console.print(f"[bold green]Starting at:[/bold green] {state['current_article']}")
        self.console.print(f"[bold red]Target:[/bold red] {state['target_article']}\n")
        
        while not state['is_complete']:
            # Display current game status
            self._display_game_status(state)
            
            # Get LLM's choice
            choice = self._get_llm_choice(state)
            self.console.print(f"\n[bold yellow]Agent chooses:[/bold yellow] {choice}")
            
            # Process the choice
            available_links = self._get_available_links(state['available_links'])
            if not available_links:
                self.console.print("[bold red]No available links to choose from![/bold red]")
                break
                
            try:
                # If choice is a number
                idx = int(choice) - 1
                if 0 <= idx < len(available_links):
                    next_article = available_links[idx]
                    self.console.print(f"[bold cyan]Moving to:[/bold cyan] {next_article}\n")
                    state, message = self.env.step(next_article)
                    if message:
                        self.console.print(f"[bold]{message}[/bold]")
                else:
                    self.console.print("[bold red]Invalid choice. Trying again.[/bold red]\n")
            except ValueError:
                self.console.print("[bold red]Invalid choice format. Trying again.[/bold red]\n")
                
        self.console.print(Panel(f"[bold green]Game completed in {state['steps_taken']} steps[/bold green]"))
        self.console.print(f"[bold]Path:[/bold] {' → '.join(state['path_taken'])}")
        return state
    
    def _display_game_status(self, state):
        """Display current game status with rich formatting"""
        # Display current article
        self.console.print(Panel(f"[bold cyan]{state['current_article']}[/bold cyan]", 
                           expand=False, 
                           border_style="cyan"))
        
        # Display article links
        self.console.print("[bold green]Available Links:[/bold green]")
        self._display_links(state['available_links'])
        
        # Display path so far
        self.console.print(f"\n[bold yellow]Steps taken:[/bold yellow] {state['steps_taken']}")
        if state['path_taken']:
            self.console.print(f"[bold yellow]Path so far:[/bold yellow] {' → '.join(state['path_taken'])}")
    
    def _display_links(self, links):
        """Display links in a nicely formatted table"""
        table = Table(show_header=False, box=SIMPLE)
        table.add_column("Number", style="dim")
        table.add_column("Link", style="green")
        table.add_column("Available", style="bold")
        
        for i, link in enumerate(links):
            # Check if link is available
            is_available = link in self.env.wiki_data
            status = "[green]✓[/green]" if is_available else "[red]✗[/red]"
            color = "green" if is_available else "red"
            table.add_row(
                f"{i+1}", 
                f"[{color}]{link}[/{color}]", 
                status
            )
            
        self.console.print(table)
    
    def _get_available_links(self, links):
        """Filter links to only those available in the wiki data"""
        return [link for link in links if link in self.env.wiki_data]
    
    @traceable(name="LLM Decision")
    def _get_llm_choice(self, state):
        """Ask LLM for next move"""
        current = state['current_article']
        target = state['target_article']
        all_links = state['available_links']
        available_links = self._get_available_links(all_links)
        path_so_far = state['path_taken']
        
        # Create prompt with relevant context (not the full article)
        prompt = f"""You are playing WikiRun, trying to navigate from one Wikipedia article to another using only links.
Current article: {current}
Target article: {target}
Available links (numbered):
{self._format_links(available_links)}

Your path so far: {' -> '.join(path_so_far)}

Think about which link is most likely to lead you toward the target article.
First, think step by step about your strategy.
Then output your choice as a number in this format: <choice>N</choice> where N is the link number.
"""
        
        # Call LLM via litellm with langsmith tracing
        response = litellm.completion(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            # metadata={
            #     "current_article": current,
            #     "target_article": target,
            #     "available_links": available_links,
            #     "steps_taken": state['steps_taken'],
            #     "path_so_far": path_so_far
            # }
        )
        
        # Extract the choice from response
        content = response.choices[0].message.content
        self.console.print(Panel(Markdown(content), title="[bold]Agent Thinking[/bold]", border_style="yellow"))
        
        # Extract choice using format <choice>N</choice>
        import re
        choice_match = re.search(r'<choice>(\d+)</choice>', content)
        if choice_match:
            return choice_match.group(1)
        else:
            # Fallback: try to find any number in the response
            numbers = re.findall(r'\d+', content)
            if numbers:
                for num in numbers:
                    if 1 <= int(num) <= len(available_links):
                        return num
            # Default to first link if no valid choice found
            return "1" if available_links else "0"
    
    def _format_links(self, links):
        """Format the list of links for the prompt"""
        return "\n".join([f"{i+1}. {link}" for i, link in enumerate(links)])

def setup_langsmith():
    """Print instructions for setting up LangSmith tracing"""
    console = Console()
    console.print(Panel("[bold yellow]LangSmith Setup Instructions[/bold yellow]"))
    console.print("To enable LangSmith tracing, set the following environment variables:")
    console.print("[bold]export LANGSMITH_API_KEY='your-api-key'[/bold]")
    console.print("Get your API key from: https://smith.langchain.com/settings")
    console.print("Once set, your WikiRun agent will log traces to your LangSmith dashboard")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        console = Console()
        console.print("[bold red]Please provide the path to Wikipedia data[/bold red]")
        console.print("Usage: python agent.py <wiki_data_path>")
        sys.exit(1)
    
    # Remind about LangSmith setup
    if not os.environ.get("LANGSMITH_API_KEY"):
        setup_langsmith()
        
    wiki_data_path = sys.argv[1]
    agent = WikiRunAgent(wiki_data_path)
    agent.run_game(start_article="Peanut", target_article="Silicon Valley")
    # agent.run_game(start_article="Silicon Valley", target_article="Peanut")