#!/usr/bin/env python3
"""
tokenviz CLI — count tokens, open web UI, pipe from stdin
"""

import sys
import os
import argparse

sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
from server import count_tokens, analyze_text, MODEL_LIMITS, main as serve_main

RESET   = "\033[0m"
BOLD    = "\033[1m"
GREEN   = "\033[32m"
YELLOW  = "\033[33m"
ORANGE  = "\033[38;5;208m"
RED     = "\033[31m"
CYAN    = "\033[36m"
MUTED   = "\033[90m"
BLUE    = "\033[34m"


def color_for_pct(pct):
    if pct >= 95: return RED
    if pct >= 80: return ORANGE
    if pct >= 60: return YELLOW
    return GREEN


def bar(pct, width=40):
    filled = int((pct / 100) * width)
    c = color_for_pct(pct)
    return c + "█" * filled + MUTED + "░" * (width - filled) + RESET


def fmt_num(n):
    return f"{n:,}"


def cmd_count(args):
    # Read text
    if args.file:
        with open(args.file, 'r', encoding='utf-8', errors='replace') as f:
            text = f.read()
    elif not sys.stdin.isatty():
        text = sys.stdin.read()
    elif args.text:
        text = args.text
    else:
        print(f"{RED}Error: provide a file, pipe text, or use --text{RESET}")
        sys.exit(1)

    model = args.model or "gpt-4o"
    data = analyze_text(text, model)
    pct = data["percent"]
    c = color_for_pct(pct)
    s = data["stats"]

    print()
    print(f"  {BOLD}tokenviz{RESET} · {MUTED}{model}{RESET}")
    print()
    print(f"  {c}{BOLD}{fmt_num(data['tokens'])}{RESET} tokens  {MUTED}/ {fmt_num(data['limit'])}{RESET}")
    print()
    print(f"  {bar(pct)}")
    print(f"  {c}{pct:.1f}%{RESET} used  ·  {GREEN}{fmt_num(data['remaining'])}{RESET} remaining")
    print()
    print(f"  {MUTED}chars    {RESET}{fmt_num(s['chars'])}")
    print(f"  {MUTED}words    {RESET}{fmt_num(s['words'])}")
    print(f"  {MUTED}lines    {RESET}{fmt_num(s['lines'])}")
    print(f"  {MUTED}tok/word {RESET}{s['avg_tokens_per_word']}")
    print()

    if data["critical"]:
        print(f"  {RED}⚠  CRITICAL: context window almost full!{RESET}")
    elif data["warning"]:
        print(f"  {YELLOW}⚠  WARNING: over 80% of context used{RESET}")

    if args.json:
        import json
        print(json.dumps(data, indent=2))


def cmd_models(_args):
    print(f"\n  {BOLD}Supported models{RESET}\n")
    for name, limit in MODEL_LIMITS.items():
        bar_len = min(30, int(30 * limit / 1_000_000))
        print(f"  {CYAN}{name:<28}{RESET} {MUTED}{fmt_num(limit):>10} tokens{RESET}")
    print()


def cmd_serve(args):
    serve_main()


def main():
    parser = argparse.ArgumentParser(
        prog='tokenviz',
        description='🔭 Real-time LLM context window visualizer'
    )
    sub = parser.add_subparsers(dest='cmd')

    # count
    p_count = sub.add_parser('count', help='Count tokens in a file or stdin')
    p_count.add_argument('file', nargs='?', help='File to analyze')
    p_count.add_argument('--text', '-t', help='Inline text to analyze')
    p_count.add_argument('--model', '-m', default='gpt-4o', help='Model to use (default: gpt-4o)')
    p_count.add_argument('--json', '-j', action='store_true', help='Output full JSON')

    # models
    sub.add_parser('models', help='List all supported models and their limits')

    # serve
    p_serve = sub.add_parser('serve', help='Start the web UI')

    if len(sys.argv) == 1:
        parser.print_help()
        print(f"\n  Quick start:\n")
        print(f"    {GREEN}tokenviz count myfile.txt{RESET}          # analyze a file")
        print(f"    {GREEN}cat prompt.txt | tokenviz count{RESET}    # pipe from stdin")
        print(f"    {GREEN}tokenviz serve{RESET}                     # open web UI")
        print(f"    {GREEN}tokenviz models{RESET}                    # list models\n")
        sys.exit(0)

    args = parser.parse_args()

    if args.cmd == 'count':
        cmd_count(args)
    elif args.cmd == 'models':
        cmd_models(args)
    elif args.cmd == 'serve':
        cmd_serve(args)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()
