#!/usr/bin/env python3
"""
PromptBridge - tmux-based split screen CLI wrapper
Top: PromptBridge input UI
Bottom: Actual CLI (opencode/claude/gemini)
"""

import os
import sys
import subprocess
import time
import signal
import shutil
from pathlib import Path
from decimal import Decimal

# prompt_toolkit for autocomplete
try:
    from prompt_toolkit import prompt
    from prompt_toolkit.completion import WordCompleter, Completer, Completion
    from prompt_toolkit.styles import Style
    from prompt_toolkit.formatted_text import HTML
    PROMPT_TOOLKIT_AVAILABLE = True
except ImportError:
    PROMPT_TOOLKIT_AVAILABLE = False
    try:
        import readline
    except ImportError:
        pass

# Package directory (installed location)
PACKAGE_DIR = Path(__file__).parent

# Working directory (where promptbridge was executed)
WORK_DIR = Path.cwd()

# Load .env (package directory or home directory)
try:
    from dotenv import load_dotenv
    # Priority: working dir > home dir > package dir
    env_locations = [
        WORK_DIR / '.env',
        Path.home() / '.promptbridge' / '.env',
        PACKAGE_DIR / '.env',
    ]
    for env_path in env_locations:
        if env_path.exists():
            load_dotenv(env_path)
            break
except ImportError:
    pass

# Auth module (all processing via server API)
try:
    from promptbridge.auth.client import AuthClient
    AUTH_AVAILABLE = True
except ImportError:
    AUTH_AVAILABLE = False
    AuthClient = None

# ANSI Colors
RESET = "\033[0m"
BOLD = "\033[1m"
CYAN = "\033[96m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
MAGENTA = "\033[95m"
DIM = "\033[2m"
BG_CYAN = "\033[46m"
BG_BLACK = "\033[40m"

SESSION_NAME = "promptbridge"
CLI_PANE = "cli"
INPUT_PANE = "input"

# Base commands (always available)
BASE_COMMANDS = {
    "!opt": "Optimize prompt and send to CLI",
    "!status": "Check login status and balance",
    "!help": "Show help",
    "!quit": "Exit PromptBridge",
}

# Commands only for logged out users
LOGGED_OUT_COMMANDS = {
    "!login": "Login",
}

# Commands only for logged in users
LOGGED_IN_COMMANDS = {
    "!logout": "Logout",
}

# Commands that require arguments (add space after selection)
COMMANDS_WITH_ARGS = {"!opt"}


def get_available_commands():
    """Return available commands based on login status"""
    commands = dict(BASE_COMMANDS)
    if AUTH_AVAILABLE and auth_client and auth_client.is_logged_in():
        commands.update(LOGGED_IN_COMMANDS)
    else:
        commands.update(LOGGED_OUT_COMMANDS)
    return commands

# Autocomplete class
if PROMPT_TOOLKIT_AVAILABLE:
    class CommandCompleter(Completer):
        def get_completions(self, document, complete_event):
            text = document.text_before_cursor
            # Only autocomplete when starting with ! and no space yet (command only)
            if text.startswith("!") and " " not in text:
                word = text
                commands = get_available_commands()
                for cmd, desc in commands.items():
                    if cmd.startswith(word):
                        # Add space for commands that require arguments
                        completion_text = cmd + " " if cmd in COMMANDS_WITH_ARGS else cmd
                        yield Completion(
                            completion_text,
                            start_position=-len(word),
                            display=f"{cmd} ->" if cmd in COMMANDS_WITH_ARGS else cmd,
                            display_meta=desc
                        )

    # Style
    PROMPT_STYLE = Style.from_dict({
        'prompt': '#00ff00 bold',
        'completion-menu.completion': 'bg:#333333 #ffffff',
        'completion-menu.completion.current': 'bg:#00aa00 #ffffff',
        'completion-menu.meta.completion': 'bg:#333333 #888888',
        'completion-menu.meta.completion.current': 'bg:#00aa00 #ffffff',
    })


# Global auth client
auth_client = None

def init_auth():
    """Initialize auth client"""
    global auth_client
    if AUTH_AVAILABLE:
        auth_client = AuthClient()


def check_tmux():
    """Check if tmux is installed"""
    result = subprocess.run(["which", "tmux"], capture_output=True)
    return result.returncode == 0


def is_inside_tmux():
    """Check if currently inside a tmux session"""
    return os.environ.get("TMUX") is not None


def kill_session():
    """Kill existing session"""
    subprocess.run(["tmux", "kill-session", "-t", SESSION_NAME],
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def do_login():
    """Perform login"""
    if not AUTH_AVAILABLE:
        print(f"{RED}Auth module is not installed.{RESET}")
        return False

    print(f"\n{CYAN}=== PromptBridge Login ==={RESET}")

    email = input(f"{GREEN}Email:{RESET} ").strip()
    if not email:
        print(f"{RED}Please enter your email.{RESET}")
        return False

    import getpass
    password = getpass.getpass(f"{GREEN}Password:{RESET} ")
    if not password:
        print(f"{RED}Please enter your password.{RESET}")
        return False

    print(f"{DIM}Logging in...{RESET}")
    success, message = auth_client.login(email, password)

    if success:
        print(f"{GREEN}✓ {message}{RESET}")
        # Update credit balance
        update_cached_balance()
        return True
    else:
        print(f"{RED}✗ {message}{RESET}")
        return False


def do_logout():
    """Perform logout"""
    if not AUTH_AVAILABLE:
        return

    success, message = auth_client.logout()
    print(f"{YELLOW}{message}{RESET}")


def show_status():
    """Show status"""
    if not AUTH_AVAILABLE:
        print(f"{RED}Auth module is not installed.{RESET}")
        return

    if not auth_client.is_logged_in():
        print(f"{YELLOW}Login required. Use !login to log in.{RESET}")
        return

    success, status = auth_client.get_status()

    if not success:
        print(f"{RED}{status.get('error', 'Failed to get status')}{RESET}")
        return

    user = status["user"]
    balance = status["balance"]
    min_balance = status["min_balance"]

    # Update credit cache
    update_cached_balance(balance)

    print(f"\n{CYAN}=== Account Status ==={RESET}")
    print(f"  {GREEN}User:{RESET} {user['username']} ({user['email']})")
    print(f"  {GREEN}Balance:{RESET} {balance} credits")
    print(f"  {DIM}Min required:{RESET} {min_balance} credits")

    balance_dec = Decimal(balance)
    min_dec = Decimal(min_balance)
    if balance_dec < min_dec:
        print(f"  {RED}⚠ Insufficient credits. Please recharge.{RESET}")
    else:
        print(f"  {GREEN}✓ Service available{RESET}")
    print()


def optimize_prompt(prompt_text: str) -> str:
    """Optimize prompt (server API call)"""

    # Check auth
    if not AUTH_AVAILABLE or not auth_client:
        print(f"{RED}Auth module is not installed.{RESET}")
        return prompt_text

    if not auth_client.is_logged_in():
        print(f"{RED}Login required. Use !login to log in.{RESET}")
        return prompt_text

    try:
        print(f"{CYAN}Optimizing...{RESET}", end="", flush=True)

        # Server API call (optimization + billing together)
        success, result = auth_client.optimize_prompt(prompt_text)

        if not success:
            print(f"\r{RED}Optimization failed: {result.get('error', 'Unknown error')}{RESET}")
            return prompt_text

        optimized = result.get("optimized_prompt", prompt_text)
        credits_deducted = result.get("credits_deducted", "0")
        new_balance = result.get("new_balance", "0")

        print(f"\r{GREEN}✓ Optimization complete!{RESET}")

        # Update balance cache
        update_cached_balance(new_balance)
        print(f"{DIM}  -> Credits: -{credits_deducted} | Balance: {new_balance}{RESET}")

        return optimized

    except Exception as e:
        print(f"\r{RED}Optimization failed: {e}{RESET}")
        return prompt_text


def send_to_cli(text: str):
    """Send text to bottom CLI panel"""
    # Escape special characters
    escaped = text.replace("\\", "\\\\").replace("'", "'\\''").replace('"', '\\"')

    # Send via tmux send-keys
    subprocess.run([
        "tmux", "send-keys", "-t", f"{SESSION_NAME}:0.1", escaped, "Enter"
    ], capture_output=True)


def clear_cli_input():
    """Clear CLI input line"""
    # Clear current line with Ctrl+U
    subprocess.run([
        "tmux", "send-keys", "-t", f"{SESSION_NAME}:0.1", "C-u"
    ], capture_output=True)


def get_terminal_size():
    """Get terminal size"""
    try:
        size = shutil.get_terminal_size()
        return size.columns, size.lines
    except:
        return 80, 24


def get_terminal_width():
    """Get terminal width"""
    return get_terminal_size()[0]


# Global credit balance cache
_cached_balance = None


def get_cached_balance():
    """Return cached credit balance"""
    global _cached_balance
    return _cached_balance


def update_cached_balance(balance=None):
    """Update credit balance"""
    global _cached_balance
    if balance is not None:
        _cached_balance = balance
    elif AUTH_AVAILABLE and auth_client and auth_client.is_logged_in():
        try:
            success, status = auth_client.get_status()
            if success:
                _cached_balance = status.get("balance", "0")
        except:
            pass
    return _cached_balance


def update_tmux_header(cli_name: str, balance: str = None):
    """Update tmux status line as header (fixed header)"""
    # Collect info
    cli_info = cli_name.upper()
    user_info = ""
    credit_info = ""

    if AUTH_AVAILABLE and auth_client:
        if auth_client.is_logged_in():
            user = auth_client.storage.get_user()
            if user:
                user_info = user['username']
            bal = balance or get_cached_balance()
            if bal:
                credit_info = f"{bal}"

    # Left: PROMPTBRIDGE | CLI name
    left_content = f"  PROMPTBRIDGE | {cli_info}  "

    # Right: user | credits or Not logged in
    if user_info and credit_info:
        right_content = f"  {user_info} | {credit_info} credits  "
    elif user_info:
        right_content = f"  {user_info}  "
    else:
        right_content = "  Not logged in  "

    # Update tmux status line
    try:
        cmds = [
            ["tmux", "set-option", "-t", SESSION_NAME, "status-left", left_content],
            ["tmux", "set-option", "-t", SESSION_NAME, "status-right", right_content],
        ]
        for cmd in cmds:
            subprocess.run(cmd, capture_output=True)
    except Exception:
        pass


def print_header(cli_name: str, clear: bool = True, balance: str = None):
    """Update header (using tmux status line)"""
    # Update fixed header via tmux status line
    update_tmux_header(cli_name, balance)

    if clear:
        # Clear screen + move cursor to first line
        print("\033[2J\033[H", end="", flush=True)
        print()  # One line margin

def print_help():
    """Print help"""
    width = get_terminal_width()
    print(f"\n{YELLOW}{'=' * width}{RESET}")
    print(f"{BOLD}PromptBridge Command Help{RESET}\n")

    commands = get_available_commands()
    for cmd, desc in commands.items():
        print(f"  {GREEN}{cmd:10}{RESET} {desc}")

    print(f"\n{DIM}Other inputs are sent directly to CLI.{RESET}")
    print(f"{YELLOW}{'=' * width}{RESET}\n")

def send_enter_to_cli():
    """Send enter to bottom CLI"""
    subprocess.run([
        "tmux", "send-keys", "-t", f"{SESSION_NAME}:0.1", "Enter"
    ], capture_output=True)


def run_input_ui(cli_name: str):
    """Input UI for top panel"""
    # Initialize auth
    init_auth()

    # Check login status and load credits
    if AUTH_AVAILABLE and auth_client:
        if auth_client.is_logged_in():
            update_cached_balance()

    print_header(cli_name)

    # Show notice only when not logged in
    if AUTH_AVAILABLE and auth_client:
        if not auth_client.is_logged_in():
            print(f"{YELLOW}! Login required. Use !login to log in.{RESET}")

    # Set up autocomplete
    completer = CommandCompleter() if PROMPT_TOOLKIT_AVAILABLE else None

    # For detecting consecutive empty inputs
    empty_input_pending = False

    while True:
        try:
            # Use prompt_toolkit (with autocomplete support)
            if PROMPT_TOOLKIT_AVAILABLE:
                raw_input = prompt(
                    HTML('<style fg="#00ff00" bold="true">[PromptBridge]&gt;</style> '),
                    completer=completer,
                    style=PROMPT_STYLE,
                    complete_while_typing=True,  # Show autocomplete on !
                )
            else:
                raw_input = input(f"{GREEN}[PromptBridge]>{RESET} ")

            # Normalize whitespace: trim + collapse multiple spaces
            user_input = ' '.join(raw_input.split())

            # Handle empty input
            if not user_input:
                if empty_input_pending:
                    # Second empty input - send enter to CLI
                    send_enter_to_cli()
                    print(f"{GREEN}✓ Sent to CLI{RESET}")
                    empty_input_pending = False
                    print_header(cli_name)
                else:
                    # First empty input
                    print(f"{DIM}Press Enter again -> Send to CLI{RESET}")
                    empty_input_pending = True
                continue

            # Reset empty input state on any input
            empty_input_pending = False

            # Command handling
            if user_input.lower() in ["!quit", "!exit"]:
                print(f"{YELLOW}Exiting...{RESET}")
                # Kill tmux session
                subprocess.run(["tmux", "kill-session", "-t", SESSION_NAME], capture_output=True)
                break

            elif user_input.lower() == "!help":
                print_help()
                input(f"{DIM}Press Enter to continue...{RESET}")
                print_header(cli_name)

            elif user_input.lower() == "!login":
                if do_login():
                    print_header(cli_name)
                else:
                    input(f"{DIM}Press Enter to continue...{RESET}")
                    print_header(cli_name)

            elif user_input.lower() == "!logout":
                do_logout()
                print_header(cli_name)

            elif user_input.lower() == "!status":
                show_status()
                input(f"{DIM}Press Enter to continue...{RESET}")
                print_header(cli_name)

            elif user_input.lower().startswith("!opt "):
                # Optimization mode
                user_prompt = user_input[5:].strip()

                if not user_prompt:
                    print(f"{RED}Please enter a prompt.{RESET}")
                    continue

                print(f"{DIM}Original: {user_prompt[:60]}{'...' if len(user_prompt) > 60 else ''}{RESET}")

                # Optimize
                optimized = optimize_prompt(user_prompt)

                print(f"{DIM}Optimized: {optimized[:80]}{'...' if len(optimized) > 80 else ''}{RESET}")

                # Send to CLI
                send_to_cli(optimized)
                print(f"{GREEN}✓ Sent!{RESET}")
                time.sleep(0.5)
                print_header(cli_name)

            else:
                # Normal input - send as is
                send_to_cli(user_input)
                print_header(cli_name)

        except KeyboardInterrupt:
            print(f"\n{YELLOW}Ctrl+C detected. Use !quit to exit.{RESET}")
        except EOFError:
            break


def select_cli():
    """CLI selection menu"""
    print("\033[2J\033[H", end="")

    print(f"""
{CYAN}{BOLD}╔═══════════════════════════════════════════════════════════╗
║                                                           ║
║ ███████╗ ██████╗   ██████╗  ███╗   ███╗ ███████╗ ████████╗║
║ ██╔══██║ ██╔══██╗ ██╔═══██╗ ████╗ ████║ ██╔══██║ ╚══██╔══╝║
║ ███████║ ██████╔╝ ██║   ██║ ██╔████╔██║ ███████║    ██║   ║
║ ██╔════╝ ██╔══██╗ ██║   ██║ ██║╚██╔╝██║ ██╔════╝    ██║   ║
║ ██║      ██║  ██║ ╚██████╔╝ ██║ ╚═╝ ██║ ██║         ██║   ║
║ ╚═╝      ╚═╝  ╚═╝  ╚═════╝  ╚═╝     ╚═╝ ╚═╝         ╚═╝   ║
║                                                           ║
║ ██████╗  ██████╗  ██╗ ██████╗   ██████╗  ███████╗         ║
║ ██╔══██╗ ██╔══██╗ ██║ ██╔══██╗ ██╔════╝  ██╔════╝         ║
║ ██████╔╝ ██████╔╝ ██║ ██║  ██║ ██║  ███╗ █████╗           ║
║ ██╔══██╗ ██╔══██╗ ██║ ██║  ██║ ██║   ██║ ██╔══╝           ║
║ ██████╔╝ ██║  ██║ ██║ ██████╔╝ ╚██████╔╝ ███████╗         ║
║ ╚═════╝  ╚═╝  ╚═╝ ╚═╝ ╚═════╝   ╚═════╝  ╚══════╝         ║
║                                                           ║
╚═══════════════════════════════════════════════════════════╝{RESET}

{BOLD}  Select a CLI:{RESET}

    {GREEN}[1]{RESET} opencode  {DIM}- Antigravity Models{RESET}
    {GREEN}[2]{RESET} claude    {DIM}- Claude Code (Anthropic){RESET}
    {GREEN}[3]{RESET} gemini    {DIM}- Gemini CLI (Google){RESET}

    {RED}[q]{RESET} Quit
""")

    while True:
        choice = input(f"{GREEN}Select [1-3, q]:{RESET} ").strip().lower()

        if choice == 'q':
            return None
        elif choice == '1':
            return "opencode"
        elif choice == '2':
            return "claude"
        elif choice == '3':
            return "gemini"
        else:
            print(f"{RED}Invalid selection.{RESET}")


def launch_tmux_session(cli_name: str, work_dir: Path = None):
    """Start tmux session"""

    # Kill existing session
    kill_session()

    # Working directory (where executed)
    if work_dir is None:
        work_dir = Path.cwd()
    work_dir_str = str(work_dir.absolute())

    # Input UI command (run as package with python -m)
    input_cmd = f"python3 -m promptbridge.main --input-mode {cli_name}"

    # Create new session + run input UI in top panel
    # -c option to specify working directory
    subprocess.run([
        "tmux", "new-session", "-d", "-s", SESSION_NAME,
        "-c", work_dir_str,  # Working directory!
        "-x", "120", "-y", "40",
        input_cmd
    ])

    # UTF-8 support settings
    subprocess.run(["tmux", "set-option", "-t", SESSION_NAME, "-g", "default-terminal", "xterm-256color"])
    subprocess.run(["tmux", "set-option", "-t", SESSION_NAME, "-g", "mouse", "on"])

    # Top fixed header (status line) - clean setup
    cmds = [
        ["tmux", "set-option", "-t", SESSION_NAME, "status", "on"],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-position", "top"],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-style", "bg=#7f19e6,fg=#ffffff,bold"],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-left", f"  PROMPTBRIDGE | {cli_name.upper()}  "],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-left-length", "40"],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-left-style", "bg=#7f19e6,fg=#ffffff,bold"],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-right", "  Not logged in  "],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-right-length", "40"],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-right-style", "bg=#7f19e6,fg=#ffffff"],
        # Clear center area completely
        ["tmux", "set-option", "-t", SESSION_NAME, "window-status-format", ""],
        ["tmux", "set-option", "-t", SESSION_NAME, "window-status-current-format", ""],
        ["tmux", "set-option", "-t", SESSION_NAME, "window-status-separator", ""],
        ["tmux", "set-option", "-t", SESSION_NAME, "status-justify", "left"],
        # Hide messages/alerts
        ["tmux", "set-option", "-t", SESSION_NAME, "display-time", "0"],
        ["tmux", "set-option", "-t", SESSION_NAME, "display-panes-time", "0"],
        # Pane border style
        ["tmux", "set-option", "-t", SESSION_NAME, "pane-border-style", "fg=#333333"],
        ["tmux", "set-option", "-t", SESSION_NAME, "pane-active-border-style", "fg=#7f19e6"],
    ]
    for cmd in cmds:
        subprocess.run(cmd, capture_output=True)

    # Create bottom panel + run CLI (same working directory)
    # 1:4 ratio (top 20%, bottom 80%)
    subprocess.run([
        "tmux", "split-window", "-t", f"{SESSION_NAME}:0",
        "-c", work_dir_str,  # Working directory!
        "-v", "-l", "80%",
        cli_name
    ])

    # Move focus to top panel (pane 0)
    subprocess.run([
        "tmux", "select-pane", "-t", f"{SESSION_NAME}:0.0"
    ])

    # Attach to session
    os.execvp("tmux", ["tmux", "attach-session", "-t", SESSION_NAME])


def main():
    import argparse

    parser = argparse.ArgumentParser(description="PromptBridge - CLI Wrapper with Prompt Optimization")
    parser.add_argument("--input-mode", help="Internal: input UI mode")
    parser.add_argument("cli", nargs="?", choices=["opencode", "claude", "gemini"], help="Specify CLI directly")

    args = parser.parse_args()

    # Input UI mode (run inside tmux)
    if args.input_mode:
        run_input_ui(args.input_mode)
        return

    # Check tmux
    if not check_tmux():
        print(f"{RED}Error: tmux is not installed.{RESET}")
        print(f"Install: brew install tmux")
        sys.exit(1)

    # Warn if already inside tmux
    if is_inside_tmux():
        print(f"{YELLOW}Warning: Already inside a tmux session.{RESET}")
        print(f"Run in a new terminal, or exit tmux first.")
        response = input("Continue anyway? [y/N]: ").strip().lower()
        if response != 'y':
            sys.exit(0)

    # Select CLI or get from argument
    if args.cli:
        cli_name = args.cli
    else:
        cli_name = select_cli()
        if not cli_name:
            print(f"{CYAN}Goodbye!{RESET}")
            sys.exit(0)

    print(f"\n{GREEN}Starting PromptBridge + {cli_name}...{RESET}")

    # Start tmux session
    launch_tmux_session(cli_name)


if __name__ == "__main__":
    main()
