import asyncio
import os
from collections import deque
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.table import Table
from rich.text import Text
from rich.align import Align
from langchain_core.messages import HumanMessage, AIMessage

from src import retriever
from src.retriever import graph
from src.react_agent.agent import COMPILED_AGENT
from src.react_agent import generate

console = Console()

async def main():
    console.print(Panel(
        Align.center("[bold magenta]Vectorless-RAG ReAct Agent Assistant[/bold magenta]\n"
                     "[dim white]Autonomous Thought -> Action -> Observation loop[/dim white]"),
        border_style="magenta"
    ))
    
    console.print("[bold yellow]Loading indices...[/bold yellow]")
    retriever.load("tree")
    console.print()
    
    console.print(Panel(
        "[bold green]Ready![/bold green] Ask a legal scenario. The agent will autonomously decide what to search.\n\n"
        "[bold white]Commands:[/bold white]\n"
        "  [cyan]exit[/cyan] / [cyan]quit[/cyan] - Close assistant\n"
        "  [cyan]clear[/cyan]       - Clear chat memory\n"
        "  [cyan]trace[/cyan]       - Toggle verbose reasoning trace (currently ON)",
        title="[bold magenta]ReAct System Status[/bold magenta]",
        border_style="magenta",
        expand=False
    ))
    console.print()
    
    MEMORY_LIMIT = 5
    history = deque(maxlen=MEMORY_LIMIT)
    trace_mode = True
    
    while True:
        try:
            query = console.input("[bold deep_sky_blue1]Query > [/bold deep_sky_blue1]")
            if query.strip().lower() in ['exit', 'quit', 'q']:
                break
            if query.strip().lower() == 'clear':
                history.clear()
                console.print("[bold green]Memory cleared![/bold green]\n")
                continue
            if query.strip().lower() == 'trace':
                trace_mode = not trace_mode
                status = "ENABLED" if trace_mode else "DISABLED"
                console.print(f"[bold yellow]Trace logs {status}.[/bold yellow]\n")
                continue
                
            if not query.strip():
                continue
                
            # If trace mode is ON, we stream updates to show the Thought-Action-Observation loop
            if trace_mode:
                console.print(f"\n[bold yellow]Agent is reasoning...[/bold yellow]")
                
                # 1. Format history messages
                messages = []
                for turn in history:
                    messages.append(HumanMessage(content=turn.get("user", "")))
                    assistant_clean = turn.get("assistant", "").split("[References]")[0].strip()
                    messages.append(AIMessage(content=assistant_clean))
                messages.append(HumanMessage(content=query))
                
                # 2. Run streaming graph
                try:
                    async for event in COMPILED_AGENT.astream(
                        {"messages": messages}, 
                        config={"recursion_limit": 10}, 
                        stream_mode="updates"
                    ):
                        for node, update in event.items():
                            if node == "agent":
                                msgs = update.get("messages", [])
                                if msgs:
                                    msg = msgs[-1]
                                    
                                    # Handle list content in agent thoughts
                                    content = msg.content
                                    if isinstance(content, list):
                                        parts = []
                                        for part in content:
                                            if isinstance(part, str):
                                                parts.append(part)
                                            elif isinstance(part, dict) and "text" in part:
                                                parts.append(part["text"])
                                            elif hasattr(part, "text"):
                                                parts.append(part.text)
                                        content = "".join(parts)
                                        
                                    if content:
                                        console.print(Panel(
                                            content.strip(), 
                                            title="[bold yellow]Agent Thought[/bold yellow]", 
                                            border_style="yellow"
                                        ))

                                    if hasattr(msg, "tool_calls") and msg.tool_calls:
                                        for tc in msg.tool_calls:
                                            console.print(f"[bold cyan]Action (Call Tool):[/bold cyan] [bold white]{tc['name']}[/bold white] with args: [magenta]{tc['args']}[/magenta]")
                            elif node == "tools":
                                msgs = update.get("messages", [])
                                if msgs:
                                    msg = msgs[-1]
                                    
                                    # Handle list content in tool observations
                                    content = msg.content
                                    if isinstance(content, list):
                                        parts = []
                                        for part in content:
                                            if isinstance(part, str):
                                                parts.append(part)
                                            elif isinstance(part, dict) and "text" in part:
                                                parts.append(part["text"])
                                            elif hasattr(part, "text"):
                                                parts.append(part.text)
                                        content = "".join(parts)
                                        
                                    preview = content[:300] + "..." if len(content) > 300 else content
                                    console.print(Panel(
                                        preview.strip(), 
                                        title="[bold green]Observation (Tool Output)[/bold green]", 
                                        border_style="green"
                                    ))
                                    console.print()
                except Exception as e:
                    console.print(f"\n[bold red]Trace Loop Error: {e}[/bold red]\n")
            
            # 3. Call standard generate interface to get final formatted answer & metadata
            with console.status("[bold yellow]Synthesizing final structured response...[/bold yellow]", spinner="dots"):
                res = await generate(
                    query=query,
                    history=list(history),
                )
            
            # 4. Print Response Metadata (ASCII only to prevent Windows console encoding crash)
            conf_badge = "[bold green][OK] ADEQUATE CONTEXT[/bold green]" if res["confidence"] > 0 else "[bold red][FAIL] INSUFFICIENT CONTEXT[/bold red]"
            border_color = "green" if res["confidence"] > 0 else "red"
            
            console.print(Panel(
                f"Status: {conf_badge}\nLatency: [cyan]{res['latency_ms']}[/cyan] ms",
                title="[bold white]Response Metadata[/bold white]",
                border_style=border_color,
                expand=False
            ))
            console.print()

            
            # 5. Print Answer using Markdown
            console.print(Panel(
                Markdown(res.get('answer', '')),
                title="[bold magenta]ReAct Final Answer[/bold magenta]",
                border_style="magenta"
            ))
            console.print()
            
            # Add to memory
            history.append({
                "user": query,
                "assistant": res.get("answer", "")
            })
            
            console.print("[dim white]" + "="*60 + "[/dim white]\n")
            
        except KeyboardInterrupt:
            break
        except Exception as e:
            console.print(f"\n[bold red]Error: {e}[/bold red]\n")

if __name__ == "__main__":
    asyncio.run(main())