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 local.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 = self.env.article_exists(link)
            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 self.env.article_exists(link)]
    
    @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")