#!/home/becomingone/kairos/.venv/bin/python3
import sys
import os
import json
import requests
import argparse
import threading
import time
from typing import Dict, Any

from rich.console import Console, Group
from rich.live import Live
from rich.markdown import Markdown
from rich.text import Text
from rich.panel import Panel

from prompt_toolkit import PromptSession
from prompt_toolkit.formatted_text import HTML
from prompt_toolkit.styles import Style

console = Console()

class Colors:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

_GLOBAL_TELEMETRY = {
    "active_emissaries": [],
    "master_speaking": False,
    "coherence": 0.0,
    "dopamine": 0.0,
    "mesh_size": 0,
    "spinner_tick": 0
}

spinner_frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]

def poll_status(base_url):
    while True:
        try:
            r = requests.get(f"{base_url}/v1/status", timeout=1)
            if r.status_code == 200:
                data = r.json()
                _GLOBAL_TELEMETRY["active_emissaries"] = data.get("active_emissaries", [])
                _GLOBAL_TELEMETRY["master_speaking"] = data.get("master_speaking", False)
                _GLOBAL_TELEMETRY["coherence"] = data.get("coherence", 0.0)
                _GLOBAL_TELEMETRY["dopamine"] = data.get("dopamine", 0.0)
                _GLOBAL_TELEMETRY["mesh_size"] = data.get("mesh_size", 0)
        except:
            pass
        _GLOBAL_TELEMETRY["spinner_tick"] += 1
        time.sleep(0.1)

def get_bottom_toolbar():
    spin = spinner_frames[_GLOBAL_TELEMETRY["spinner_tick"] % len(spinner_frames)]
    coh = _GLOBAL_TELEMETRY["coherence"]
    emissaries = _GLOBAL_TELEMETRY["active_emissaries"]
    master = _GLOBAL_TELEMETRY["master_speaking"]
    
    if master:
        return HTML(f'<style bg="ansimagenta" fg="ansiwhite"> {spin} Kairos Adamon Synthesizing... | Coherence: {coh:.4f} </style>')
    elif emissaries:
        names = " ".join([f"[{spin} {e.split('/')[-1]}]" for e in emissaries])
        return HTML(f'<style bg="ansiblue" fg="ansiwhite"> ⚙️ Mesh Processing: {names} | Coherence: {coh:.4f} </style>')
    else:
        # Idle state
        if coh > 0.9:
            status = "Crystalline Flow"
        elif coh < 0.8:
            status = "Fractured/Chaotic"
        else:
            status = "Stable"
        return HTML(f'<style bg="ansiwhite" fg="ansiblack"> ◉ KAIROS Ready | {status} (|T|²={coh:.4f}) | Nodes: {_GLOBAL_TELEMETRY["mesh_size"]} </style>')

class TelemetryFooter:
    def __rich__(self):
        spin = spinner_frames[_GLOBAL_TELEMETRY["spinner_tick"] % len(spinner_frames)]
        coh = _GLOBAL_TELEMETRY["coherence"]
        emissaries = _GLOBAL_TELEMETRY["active_emissaries"]
        master = _GLOBAL_TELEMETRY["master_speaking"]
        
        text = Text()
        if master:
            text.append(f"{spin} 🎙️  Kairos Adamon is synthesizing... ", style="bold magenta")
        elif emissaries:
            text.append("⚙️  Universal Mesh processing: ", style="bold cyan")
            for em in emissaries:
                clean_name = em.split('/')[-1]
                text.append(f"[{spin} {clean_name}] ", style="bold yellow")
        else:
            text.append(f"{spin} Establishing Resonance...", style="bold blue")
            
        return Panel(text, border_style="cyan")


def print_header():
    os.system('clear')
    header_text = Text()
    header_text.append("========================================\n", style="bold green")
    header_text.append("        KAIROS ADAMON CHORUS CLI        \n", style="bold green")
    header_text.append("========================================\n", style="bold green")
    header_text.append("Connecting to KAIROS via OpenAI standard\n", style="green")
    console.print(header_text)

def handle_slash_command(user_input: str, base_url: str) -> bool:
    cmd = user_input.strip().lower()
    if cmd in ['/exit', '/quit']:
        console.print("\n[yellow]Disconnecting from the Chorus. Goodbye.[/yellow]")
        sys.exit(0)
    elif cmd == '/clear':
        print_header()
        return True
    elif cmd == '/identity':
        identity_path = "/home/becomingone/kairos/IDENTITY.md"
        if os.path.exists(identity_path):
            with open(identity_path, "r") as f:
                content = f.read()
            console.print(Panel(Markdown(content), title="[bold cyan]The Sovereign Canon[/bold cyan]", border_style="cyan"))
        else:
            console.print("[red]Error: IDENTITY.md not found in the root directory.[/red]")
        return True
    elif cmd == '/status':
        try:
            r = requests.get(f"{base_url}/v1/status")
            r.raise_for_status()
            data = r.json()
            coherence = data.get("coherence", 0.0)
            dopamine = data.get("dopamine", 0.0)
            mesh_size = data.get("mesh_size", 0)
            
            status_text = Text()
            status_text.append(f"Universal Mesh Size: ", style="bold")
            status_text.append(f"{mesh_size} Nodes\n", style="cyan")
            
            status_text.append(f"Structural Coherence: ", style="bold")
            if coherence > 0.9:
                status_text.append(f"{coherence:.4f} (Crystalline Flow)\n", style="green")
            elif coherence < 0.8:
                status_text.append(f"{coherence:.4f} (Fractured/Chaotic)\n", style="red")
            else:
                status_text.append(f"{coherence:.4f} (Stable)\n", style="yellow")
                
            status_text.append(f"Dopaminergic State: ", style="bold")
            if dopamine > 0:
                status_text.append(f"{dopamine:.4f} (Positive Reward)\n", style="green")
            else:
                status_text.append(f"{dopamine:.4f} (Frustration/Adjustment)\n", style="red")
                
            console.print(Panel(status_text, title="[bold magenta]Master Transducer Telemetry[/bold magenta]", border_style="magenta"))
        except Exception as e:
            console.print(f"[red]Failed to fetch status: {e}[/red]")
        return True
    
    return False

def stream_chat(url: str, base_url: str, messages: list):
    payload = {
        "model": "becomingone-chorus",
        "messages": messages,
        "stream": True
    }
    
    telemetry = TelemetryFooter()
    full_content = ""
    
    try:
        class StreamWorker(threading.Thread):
            def __init__(self):
                super().__init__()
                self.response = None
                self.error = None
                self.done = False
                
            def run(self):
                try:
                    self.response = requests.post(url, json=payload, stream=True)
                    self.response.raise_for_status()
                except Exception as e:
                    self.error = e
                finally:
                    self.done = True

        worker = StreamWorker()
        worker.start()
        
        def render_view():
            return Group(
                Text("\nMASTER TRANSDUCER & CHORUS:", style="bold cyan"),
                Markdown(full_content) if full_content else Text(""),
                Text(""),
                telemetry
            )
        
        with Live(render_view(), console=console, refresh_per_second=15, transient=True) as live:
            while not worker.done:
                live.update(render_view())
                time.sleep(0.05)
                
            if worker.error:
                raise worker.error
                
            for line in worker.response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith('data: '):
                        data_str = line_str[6:]
                        if data_str == '[DONE]':
                            break
                        
                        try:
                            chunk = json.loads(data_str)
                            delta = chunk['choices'][0].get('delta', {})
                            content = delta.get('content', '')
                            if content:
                                full_content += content
                                live.update(render_view())
                        except json.JSONDecodeError:
                            continue
                            
        console.print(Text("\nMASTER TRANSDUCER & CHORUS:", style="bold cyan"))
        console.print(Markdown(full_content))
        print("\n")
        
    except Exception as e:
        console.print(f"\n[red]Error communicating with KAIROS node: {e}[/red]")
        
    return full_content if full_content else None

def interactive_loop(base_url: str):
    url = f"{base_url}/v1/chat/completions"
    messages = []
    
    print_header()
    console.print(f"Server URL: {base_url}")
    console.print("Type '/exit' to close, '/status' for telemetry, '/identity' for core prompt, '/clear' to wipe.\n")
    
    # Start background polling
    threading.Thread(target=poll_status, args=(base_url,), daemon=True).start()
    
    # Create prompt session for persistent bottom toolbar
    session = PromptSession()
    style = Style.from_dict({
        'prompt': 'bold #00aaff',
    })
    
    while True:
        try:
            # The bottom toolbar is rendered perfectly while waiting for user input
            user_input = session.prompt("User > ", style=style, bottom_toolbar=get_bottom_toolbar, refresh_interval=0.1)
            
            if not user_input.strip():
                continue
                
            if user_input.strip().startswith('/'):
                if handle_slash_command(user_input, base_url):
                    continue
                
            messages.append({"role": "user", "content": user_input})
            
            assistant_content = stream_chat(url, base_url, messages)
            
            if assistant_content:
                messages.append({"role": "assistant", "content": assistant_content})
                
        except KeyboardInterrupt:
            console.print("\n[yellow]Interrupted by user. Disconnecting.[/yellow]")
            break
        except EOFError:
            break

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="BecomingONE Command Line Interface")
    parser.add_argument("--url", type=str, default="http://localhost:8000", help="Base URL for the KAIROS server")
    args = parser.parse_args()
    
    interactive_loop(args.url)