import sys
import os
import json

# --- ANSI Styling Palette ---
CYAN = "\033[96m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
CLEAR = "\033[2J\033[H"

# --- ANSI Layout Sequences ---
SAVE_CURSOR = "\033[s"
RESTORE_CURSOR = "\033[u"
CLEAR_LINE = "\033[K"

class TerminalUIManager:
    def __init__(self):
        self.term_width, self.term_height = os.get_terminal_size()
        self.current_persona = "Orchestrator"
        self.active_tool = None
        self.token_usage = 0
        self.max_tokens = 16384
        self.is_thinking = False
        self.has_set_region = False
    
    def setup_screen(self):
        """Initializes the terminal layout with a scrolling region."""
        sys.stdout.write(CLEAR)
        sys.stdout.flush()
        # Set scrolling region from row 2 to (height - 1)
        self.term_width, self.term_height = os.get_terminal_size()
        sys.stdout.write(f"\033[2;{self.term_height - 1}r")
        # Move cursor to row 2
        sys.stdout.write("\033[2;1H")
        sys.stdout.flush()
        self.has_set_region = True
        self.redraw_header()
        self.redraw_footer()

    def cleanup(self):
        """Restores the terminal scrolling region."""
        if self.has_set_region:
            sys.stdout.write("\033[r") # Reset scrolling region
            sys.stdout.write(f"\033[{self.term_height};1H")
            sys.stdout.write("\n")
            sys.stdout.flush()

    def update_state(self, persona=None, active_tool=None, token_usage=None):
        """Update internal state and redraw header/footer if necessary."""
        changed_header = False
        changed_footer = False
        
        if persona is not None and persona != self.current_persona:
            self.current_persona = persona
            changed_header = True
            
        if active_tool is not None and active_tool != self.active_tool:
            self.active_tool = active_tool
            changed_footer = True
            
        if token_usage is not None and token_usage != self.token_usage:
            self.token_usage = token_usage
            changed_footer = True
            
        if changed_header:
            self.redraw_header()
        if changed_footer:
            self.redraw_footer()

    def redraw_header(self):
        """Draws the fixed top header."""
        sys.stdout.write(SAVE_CURSOR)
        sys.stdout.write("\033[1;1H") # Move to top left
        
        persona_name = self.current_persona.split('/')[-1].capitalize()
        header_text = f"⚙ SmarterAgents │ Persona: {persona_name}"
        padding = max(0, self.term_width - len(header_text) - 1)
        
        sys.stdout.write(f"{CYAN}{BOLD}{header_text} {'─' * padding}{RESET}{CLEAR_LINE}")
        sys.stdout.write(RESTORE_CURSOR)
        sys.stdout.flush()

    def redraw_footer(self):
        """Draws the fixed bottom footer."""
        sys.stdout.write(SAVE_CURSOR)
        sys.stdout.write(f"\033[{self.term_height};1H") # Move to bottom left
        
        # Calculate budget bar
        pct = min(1.0, self.token_usage / self.max_tokens) if self.max_tokens > 0 else 0
        filled = int(pct * 10)
        empty = 10 - filled
        bar = "█" * filled + "░" * empty
        pct_str = f"{int(pct * 100)}%"
        
        status_text = "Idle"
        if self.active_tool:
            status_text = f"⚡ Running Tool: {self.active_tool}"
            
        footer_text = f" Token Budget: [{bar}] {pct_str} │ Status: {status_text} "
        padding = max(0, self.term_width - len(footer_text))
        
        sys.stdout.write(f"{DIM}{footer_text}{' ' * padding}{RESET}{CLEAR_LINE}")
        sys.stdout.write(RESTORE_CURSOR)
        sys.stdout.flush()

    def start_thought(self):
        """Starts rendering a thought block."""
        if not self.is_thinking:
            self.is_thinking = True
            sys.stdout.write(f"\n{CYAN}{DIM}[💭 Agent Thinking...]\n")
            sys.stdout.flush()

    def end_thought(self):
        """Ends rendering a thought block."""
        if self.is_thinking:
            self.is_thinking = False
            sys.stdout.write(f"{RESET}\n{CYAN}{DIM}[✓ Thought Complete]{RESET}\n")
            sys.stdout.flush()

    def stream_thought_chunk(self, chunk):
        """Streams raw thought characters."""
        sys.stdout.write(chunk)
        sys.stdout.flush()

    def print_message(self, message, role="Assistant"):
        """Prints a standard message."""
        color = GREEN if role == "Assistant" else YELLOW
        sys.stdout.write(f"\n{color}{BOLD}{role} >{RESET} {message}\n")
        sys.stdout.flush()

    def log_tool_call(self, tool_name, args):
        """Prints a nicely formatted tool invocation box."""
        sys.stdout.write(f"\n{YELLOW}{BOLD}╭─[🔧 Tool Execution: {tool_name}]" + "─" * max(0, self.term_width - 25 - len(tool_name)) + f"{RESET}\n")
        try:
            formatted_args = json.dumps(args, indent=2)
            for line in formatted_args.splitlines():
                sys.stdout.write(f"{YELLOW}│{RESET} {line}\n")
        except:
            sys.stdout.write(f"{YELLOW}│{RESET} {args}\n")
        sys.stdout.write(f"{YELLOW}╰" + "─" * (self.term_width - 1) + f"{RESET}\n")
        sys.stdout.flush()

    def print_system_info(self, message, is_error=False):
        color = RED if is_error else CYAN
        sys.stdout.write(f"\n{color}{BOLD}[System]{RESET} {message}\n")
        sys.stdout.flush()