#!/usr/bin/env python3
"""Terminal CLI -- AhoosAI in the terminal. Terminal CLI made by AhoosAI (https://ahoos-ai.site).

    terminal-cli                   a conversation, in this terminal
    terminal-cli "a question"      one answer, then back to the shell
    cat notes.txt | terminal-cli "sum it up"
    terminal-cli login             keep an API key (from https://ahoos-ai.site/services/platforme/)

It talks to the AhoosAI API (Nimbus 2 Apex) with one key, and does what the online studio does:
the reasoning as it comes, Apex's two dials (thinking 1-20, temperature 1-10), web search, plan
mode, files and images in a message, chats kept and reopened, a context meter and compaction,
custom instructions, and answers drawn for a terminal -- headings, code, tables, charts, notes.
One thing the online studio cannot: agent mode (/agent) lets the model read and write files and
run commands in the folder you are in, each change asked first.

Only Python's standard library: run it as `python terminal_cli.py`, or as the built `terminal-cli` program.
"""

from __future__ import annotations

import argparse
import base64
import difflib
import json
import mimetypes
import os
import re
import shutil
import subprocess
import sys
import textwrap
import threading
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path

VERSION = "1.0.0"
NAME = "terminal-cli"
TITLE = "Terminal CLI"
API = os.environ.get("TERMINAL_CLI_API_URL") or os.environ.get("AHOOS_API_URL", "https://ahoos-ai.site/services/api/v1").rstrip("/")
KEYS_PAGE = "https://ahoos-ai.site/services/platforme/"
HOME = Path(os.environ.get("TERMINAL_CLI_HOME") or os.environ.get("AHOOS_HOME") or Path.home() / ".terminal-cli")
CONFIG = HOME / "config.json"
CHATS = HOME / "chats"
CONTEXT = 40960                                    # the engine's slot, in tokens
MAX_FILE = 200_000                                 # characters of one attached file
IMAGE_BYTES = 6 * 1024 * 1024

# -- words -----------------------------------------------------------------------------------------------

WORDS = {
    "hello": ("made by AhoosAI · Nimbus 2 Apex · /help for commands", "made by AhoosAI · نیمبوس ۲ ایپکس · /help برای فرمان‌ها"),
    "no_key": ("No API key yet. Make one at {page} and run: terminal-cli login", "هنوز کلید API ندارید. در {page} بسازید و اجرا کنید: terminal-cli login"),
    "key_prompt": ("API key (nim_...): ", "کلید API (nim_...): "),
    "key_saved": ("Key saved in {path}.", "کلید در {path} ذخیره شد."),
    "key_bad": ("That does not look like an AhoosAI key (nim_...).", "این کلید اهوس‌ای‌آی نیست (nim_...)."),
    "waking": ("waking the model -- the first answer after a quiet spell takes about a minute", "مدل در حال روشن شدن است؛ پاسخ اول بعد از مدتی بی‌کاری حدود یک دقیقه طول می‌کشد"),
    "thinking": ("thinking", "در حال فکر"),
    "thought": ("thought {n} words", "{n} کلمه فکر"),
    "searching": ("searching", "جست‌وجو"),
    "reading": ("reading", "خواندن"),
    "results": ("{n} results", "{n} نتیجه"),
    "stopped": ("stopped", "متوقف شد"),
    "lost": ("The connection was lost before the answer ended. /retry asks again.", "اتصال پیش از پایان پاسخ قطع شد. ‏/retry دوباره می‌پرسد."),
    "context": ("context {used} of {limit} tokens ({pct}%)", "حافظه‌ی گفتگو {used} از {limit} توکن ({pct}٪)"),
    "compacting": ("compacting the conversation...", "فشرده‌سازی گفتگو..."),
    "compacted": ("compacted: the summary stands in for {n} messages", "فشرده شد: خلاصه جای {n} پیام را گرفت"),
    "new_chat": ("new chat", "گفتگوی تازه"),
    "attached": ("attached {name} ({size})", "پیوست شد: {name} ({size})"),
    "image_used": ("Used apex-imagexo-v1", "Used apex-imagexo-v1"),
    "ask_tool": ("{what}  allow? [y]es / [n]o / [a]lways this turn: ", "{what}  اجازه؟ [y] بله / [n] نه / [a] همیشه در این نوبت: "),
    "agent_on": ("agent mode: the model may read and write files and run commands in {cwd} (each change asked first)",
                 "حالت ایجنت: مدل می‌تواند در {cwd} فایل بخواند و بنویسد و فرمان اجرا کند (هر تغییر پیش از انجام پرسیده می‌شود)"),
    "on": ("on", "روشن"), "off": ("off", "خاموش"),
    "bye": ("bye", "خداحافظ"),
}
LANG = "en"


def w(key: str, **kw) -> str:
    en, fa = WORDS[key]
    return (fa if LANG == "fa" else en).format(**kw)


# -- the terminal ----------------------------------------------------------------------------------------

def _ansi_ok() -> bool:
    if os.environ.get("NO_COLOR") or not sys.stdout.isatty():
        return False
    if os.name == "nt":
        try:
            import ctypes

            kernel = ctypes.windll.kernel32
            handle = kernel.GetStdHandle(-11)
            mode = ctypes.c_uint32()
            if kernel.GetConsoleMode(handle, ctypes.byref(mode)):
                kernel.SetConsoleMode(handle, mode.value | 0x0004)
        except Exception:                                       # noqa: BLE001 - no colour is fine
            return False
    return True


COLOUR = _ansi_ok()
if hasattr(sys.stdout, "reconfigure"):
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        sys.stdin.reconfigure(encoding="utf-8", errors="replace")
    except (OSError, ValueError):
        pass


def c(text: str, code: str) -> str:
    return f"\033[{code}m{text}\033[0m" if COLOUR and text else text


def dim(t): return c(t, "2")
def bold(t): return c(t, "1")
def blue(t): return c(t, "38;5;69")
def green(t): return c(t, "32")
def yellow(t): return c(t, "33")
def red(t): return c(t, "31")
def cyan(t): return c(t, "36")


def cols() -> int:
    return max(40, min(110, shutil.get_terminal_size((90, 24)).columns))


def say(line: str = "") -> None:
    print(line, flush=True)


# -- answers drawn for a terminal ------------------------------------------------------------------------

def inline(text: str) -> str:
    codes: list[str] = []
    text = re.sub(r"`([^`\n]+)`", lambda m: codes.append(m.group(1)) or f"\0{len(codes) - 1}\0", text)
    text = re.sub(r"\*\*([^*\n]+)\*\*|__([^_\n]+)__", lambda m: bold(m.group(1) or m.group(2)), text)
    text = re.sub(r"(?<![*\w])\*([^*\n]+)\*(?!\w)", lambda m: c(m.group(1), "3"), text)
    text = re.sub(r"\[([^\]]+)\]\((https?://[^)\s]+)\)", lambda m: f"{m.group(1)} {dim('(' + m.group(2) + ')')}", text)
    return re.sub(r"\0(\d+)\0", lambda m: cyan(codes[int(m.group(1))]), text)


def visible(text: str) -> int:
    return len(re.sub(r"\033\[[0-9;]*m", "", text))


def draw_table(rows: list[str]) -> list[str]:
    cells = [[x.strip() for x in r.strip().strip("|").split("|")] for r in rows]
    if len(cells) > 1 and all(re.fullmatch(r":?-{2,}:?", x or "--") for x in cells[1]):
        cells.pop(1)
    n = max(len(r) for r in cells)
    cells = [r + [""] * (n - len(r)) for r in cells]
    shown = [[inline(x) for x in r] for r in cells]
    widths = [max(visible(r[i]) for r in shown) for i in range(n)]
    room = cols() - 3 * n - 1
    while sum(widths) > room and max(widths) > 8:
        widths[widths.index(max(widths))] -= 1
    def cut(s: str, width: int) -> str:
        if visible(s) <= width:
            return s + " " * (width - visible(s))
        plain = re.sub(r"\033\[[0-9;]*m", "", s)
        return plain[:width - 1] + "…"
    line = lambda l, m, r: dim(l + m.join("─" * (x + 2) for x in widths) + r)
    out = [line("┌", "┬", "┐")]
    for i, r in enumerate(shown):
        out.append(dim("│") + dim("│").join(" " + (bold(cut(x, widths[j])) if i == 0 else cut(x, widths[j])) + " "
                                         for j, x in enumerate(r)) + dim("│"))
        if i == 0:
            out.append(line("├", "┼", "┤"))
    out.append(line("└", "┴", "┘"))
    return out


def draw_chart(code: str) -> list[str] | None:
    """A ```chart block (the studio's JSON) as bars of blocks; None if it is not one."""
    try:
        spec = json.loads(code)
    except ValueError:
        return None
    if isinstance(spec, list):
        spec = {"data": spec}
    labels, series = spec.get("labels") or [], spec.get("series")
    if not series:
        data = spec.get("data") or spec.get("values") or []
        if isinstance(data, dict):
            labels, data = list(data), list(data.values())
        if data and isinstance(data[0], list):
            labels, data = [d[0] for d in data], [d[1] for d in data]
        series = [{"name": spec.get("name", ""), "data": data}]
    try:
        series = [{"name": str(s.get("name", "")), "data": [float(v) for v in s.get("data", [])]} for s in series]
    except (TypeError, ValueError, AttributeError):
        return None
    if not series or not series[0]["data"]:
        return None
    labels = [str(x) for x in labels] or [str(i + 1) for i in range(len(series[0]["data"]))]
    top = max(max(s["data"] or [0]) for s in series) or 1
    label_w = min(22, max(len(x) for x in labels))
    bar_w = max(10, cols() - label_w - 16)
    paints = [blue, yellow, green, cyan, red]
    out = [bold(str(spec.get("title", ""))) if spec.get("title") else ""]
    total = sum(v for v in series[0]["data"] if v > 0) or 1
    for i, label in enumerate(labels):
        for j, s in enumerate(series):
            if i >= len(s["data"]):
                continue
            v = s["data"][i]
            n = max(0, int(round(bar_w * v / top)))
            name = label[:label_w].ljust(label_w) if j == 0 else " " * label_w
            extra = f" ({100 * v / total:.0f}%)" if str(spec.get("type", "")).lower() in ("pie", "donut", "doughnut") else ""
            out.append(f"{name} {paints[j % len(paints)]('█' * n or '▏')} {v:g}{(' ' + spec['unit']) if spec.get('unit') else ''}{extra}")
    if len(series) > 1:
        out.append("  ".join(paints[j % len(paints)]("█") + f" {s['name']}" for j, s in enumerate(series)))
    return [x for x in out if x != ""]


PART = {"note": ("i", blue), "info": ("i", blue), "tip": ("✓", green), "success": ("✓", green),
        "warning": ("!", yellow), "caution": ("!", yellow), "important": ("!", yellow), "danger": ("×", red)}


class Renderer:
    """Markdown as it arrives, a line at a time; tables, charts and code blocks once they close."""

    def __init__(self) -> None:
        self.buffer = ""
        self.fence: str | None = None
        self.fence_lang = ""
        self.fence_lines: list[str] = []
        self.table: list[str] = []
        self.depth = 0

    def feed(self, text: str) -> None:
        self.buffer += text
        while "\n" in self.buffer:
            line, self.buffer = self.buffer.split("\n", 1)
            self.line(line)

    def end(self) -> None:
        if self.buffer:
            self.line(self.buffer)
            self.buffer = ""
        self.close_table()
        if self.fence is not None:
            self.close_fence()

    def out(self, text: str) -> None:
        pad = dim("│ ") * self.depth
        width = cols() - 2 * self.depth
        if visible(text) <= width:
            print(pad + text, flush=True)
            return
        for part in textwrap.wrap(text, width=width, break_long_words=False, replace_whitespace=False) or [""]:
            print(pad + part, flush=True)

    def close_table(self) -> None:
        if self.table:
            for row in draw_table(self.table):
                print(dim("│ ") * self.depth + row, flush=True)
            self.table = []

    def close_fence(self) -> None:
        code = "\n".join(self.fence_lines)
        drawn = draw_chart(code) if self.fence_lang in ("chart", "graph", "json", "") else None
        if drawn and (self.fence_lang in ("chart", "graph") or re.search(r'"type"\s*:\s*"(bar|line|pie|donut|hbar|area)"', code)):
            for row in drawn:
                self.out(row)
        else:
            label = self.fence_lang or "text"
            print(dim("  ┌─ " + label), flush=True)
            for row in self.fence_lines:
                print(dim("  │ ") + cyan(row), flush=True)
            print(dim("  └─"), flush=True)
        self.fence, self.fence_lines, self.fence_lang = None, [], ""

    def line(self, line: str) -> None:
        if self.fence is not None:
            if line.strip().startswith(self.fence):
                self.close_fence()
            else:
                self.fence_lines.append(line)
            return
        opened = re.match(r"\s*(`{3,}|~{3,})\s*([\w+#.-]*)", line)
        if opened:
            self.close_table()
            self.fence, self.fence_lang = opened.group(1), opened.group(2).lower()
            return
        if re.match(r"\s*\|.*\|\s*$", line):
            self.table.append(line)
            return
        self.close_table()
        part = re.match(r"\s*:::\s*([A-Za-z]+)\s*(.*)$", line)
        if part:
            kind, title = part.group(1).lower(), part.group(2).strip()
            if kind in ("vertical", "down"):
                return
            mark, paint = PART.get(kind, ("§", bold))
            heading = title or {"details": "More", "more": "More", "steps": "Steps", "section": ""}.get(kind, kind.title())
            self.out(paint(f"{mark} {heading}" if heading else mark) if kind not in ("section", "part") else bold("■ " + heading))
            self.depth += 1
            return
        if re.match(r"\s*:::\s*$", line):
            self.depth = max(0, self.depth - 1)
            return
        heading = re.match(r"(#{1,6})\s+(.*)", line)
        if heading:
            self.out("")
            self.out(bold(blue(inline(heading.group(2)))) if len(heading.group(1)) <= 2 else bold(inline(heading.group(2))))
            return
        if re.match(r"\s*([-*_])(\s*\1){2,}\s*$", line):
            self.out(dim("─" * min(40, cols() - 4)))
            return
        item = re.match(r"(\s*)([-*+]|\d+[.)])\s+(.*)", line)
        if item:
            bullet = "•" if item.group(2) in "-*+" else item.group(2)
            task = re.match(r"\[([ xX])\]\s+(.*)", item.group(3))
            if task:
                bullet, rest = ("☑" if task.group(1) != " " else "☐"), task.group(2)
            else:
                rest = item.group(3)
            self.out(item.group(1) + blue(bullet) + " " + inline(rest))
            return
        quote = re.match(r"\s*>\s?(.*)", line)
        if quote:
            self.out(dim("▌ ") + c(inline(quote.group(1)), "3"))
            return
        self.out(inline(line))


# -- settings, chats -------------------------------------------------------------------------------------

def load_config() -> dict:
    try:
        return json.loads(CONFIG.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}


def save_config(config: dict) -> None:
    HOME.mkdir(parents=True, exist_ok=True)
    CONFIG.write_text(json.dumps(config, ensure_ascii=False, indent=1), encoding="utf-8")
    try:
        os.chmod(CONFIG, 0o600)
    except OSError:
        pass


def api_key(config: dict) -> str:
    return os.environ.get("TERMINAL_CLI_API_KEY") or os.environ.get("AHOOS_API_KEY") or config.get("key", "")


class Chat:
    def __init__(self, data: dict | None = None) -> None:
        data = data or {}
        self.id = data.get("id") or time.strftime("%Y%m%d-%H%M%S-") + uuid.uuid4().hex[:4]
        self.title = data.get("title", "")
        self.messages: list[dict] = data.get("messages", [])
        self.context = data.get("context") or {}

    @property
    def path(self) -> Path:
        return CHATS / f"{self.id}.json"

    def save(self) -> None:
        if not self.messages:
            return
        CHATS.mkdir(parents=True, exist_ok=True)
        self.path.write_text(json.dumps({"id": self.id, "title": self.title, "messages": self.messages,
                                         "context": self.context, "updated": time.time()}, ensure_ascii=False, indent=1),
                             encoding="utf-8")

    @staticmethod
    def listing() -> list[dict]:
        out = []
        for f in sorted(CHATS.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True) if CHATS.is_dir() else []:
            try:
                d = json.loads(f.read_text(encoding="utf-8"))
                out.append({"id": d["id"], "title": d.get("title") or "?", "n": len(d.get("messages", [])), "at": f.stat().st_mtime})
            except (OSError, ValueError, KeyError):
                continue
        return out

    def history(self, memory: int = 12) -> list[dict]:
        """What the model is sent: a compaction's summary in place of what it covers, then the last turns."""
        start, summary = 0, None
        for i, m in enumerate(self.messages):
            if m.get("summary"):
                start, summary = i + 1, m
        turns = [m for m in self.messages[start:] if m.get("role") in ("user", "assistant") and not m.get("summary")]
        turns = turns[-memory * 2:]
        out = []
        if summary:
            out += [{"role": "user", "content": "Summary of our conversation so far, in place of the earlier messages:\n\n"
                                                + summary["content"]}, {"role": "assistant", "content": "Understood."}]
        for m in turns:
            out.append({"role": m["role"], "content": m.get("sent") or m["content"]})
        return out


# -- the API ---------------------------------------------------------------------------------------------

class ApiError(RuntimeError):
    pass


def request(path: str, key: str, body: dict | None = None, timeout: float = 30):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(API + path, data=data, method="POST" if data else "GET", headers={
        "Authorization": f"Bearer {key}", "Content-Type": "application/json", "User-Agent": f"terminal-cli/{VERSION}",
        "Accept": "text/event-stream" if body and body.get("stream") else "application/json"})
    try:
        return urllib.request.urlopen(req, timeout=timeout)
    except urllib.error.HTTPError as error:
        detail = error.read(800).decode("utf-8", "replace")
        try:
            detail = json.loads(detail).get("error", {}).get("message") or detail
        except (ValueError, AttributeError):
            pass
        raise ApiError(f"{error.code}: {str(detail)[:300]}") from error
    except (urllib.error.URLError, OSError) as error:
        raise ApiError(f"could not reach {API}: {getattr(error, 'reason', error)}") from error


def events(response):
    """The SSE stream's JSON chunks (keepalive comments are skipped)."""
    for raw in response:
        line = raw.decode("utf-8", "replace").strip()
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if data == "[DONE]":
            return
        try:
            yield json.loads(data)
        except ValueError:
            continue


# -- agent mode: the model's hands on this computer ------------------------------------------------------

TOOLS = [
    {"type": "function", "function": {"name": "list_files", "description": "List the files under a folder of the project (relative path; '.' is the project).",
                                      "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}}},
    {"type": "function", "function": {"name": "read_file", "description": "Read a text file of the project.",
                                      "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}}},
    {"type": "function", "function": {"name": "write_file", "description": "Create or overwrite a file with its complete content.",
                                      "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
                                                     "required": ["path", "content"]}}},
    {"type": "function", "function": {"name": "edit_file", "description": "Replace one exact piece of a file's text with another.",
                                      "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "old": {"type": "string"},
                                                                                      "new": {"type": "string"}}, "required": ["path", "old", "new"]}}},
    {"type": "function", "function": {"name": "search_files", "description": "Find lines matching a regular expression in the project's text files.",
                                      "parameters": {"type": "object", "properties": {"pattern": {"type": "string"}, "path": {"type": "string"}},
                                                     "required": ["pattern"]}}},
    {"type": "function", "function": {"name": "run_command", "description": "Run a shell command in the project folder (build, test, install, start) and see its output.",
                                      "parameters": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}}},
]
SKIP = {".git", "node_modules", "__pycache__", ".venv", "venv", "dist", "build", ".idea", ".vscode"}


class Hands:
    def __init__(self, root: Path) -> None:
        self.root = root.resolve()
        self.always = False

    def inside(self, rel: str) -> Path:
        p = (self.root / (rel or ".")).resolve()
        if p != self.root and self.root not in p.parents:
            raise ValueError(f"{rel} is outside the project folder")
        return p

    def ask(self, what: str) -> bool:
        if self.always:
            return True
        try:
            answer = input(yellow(w("ask_tool", what=what))).strip().lower()
        except EOFError:
            return False
        if answer in ("a", "always", "همیشه"):
            self.always = True
            return True
        return answer in ("y", "yes", "بله", "ب")

    def run(self, name: str, args: dict) -> str:
        try:
            if name == "list_files":
                base = self.inside(args.get("path", "."))
                found = []
                for p in sorted(base.rglob("*")):
                    if any(part in SKIP for part in p.relative_to(self.root).parts):
                        continue
                    if p.is_file():
                        found.append(str(p.relative_to(self.root)).replace("\\", "/"))
                    if len(found) >= 300:
                        break
                return "\n".join(found) or "(empty)"
            if name == "read_file":
                p = self.inside(args["path"])
                text = p.read_text(encoding="utf-8", errors="replace")
                return text[:MAX_FILE] + ("\n[... cut]" if len(text) > MAX_FILE else "")
            if name == "search_files":
                rx = re.compile(args["pattern"])
                hits = []
                for p in self.inside(args.get("path", ".")).rglob("*"):
                    if p.is_file() and not any(part in SKIP for part in p.relative_to(self.root).parts) and p.stat().st_size < 1_000_000:
                        for i, line in enumerate(p.read_text(encoding="utf-8", errors="ignore").splitlines(), 1):
                            if rx.search(line):
                                hits.append(f"{p.relative_to(self.root)}:{i}: {line.strip()[:200]}")
                                if len(hits) >= 100:
                                    return "\n".join(hits)
                return "\n".join(hits) or "no match"
            if name == "write_file":
                p = self.inside(args["path"])
                old = p.read_text(encoding="utf-8", errors="replace") if p.exists() else ""
                self.show_diff(args["path"], old, args["content"])
                if not self.ask(f"write {args['path']} ({len(args['content'].splitlines())} lines)"):
                    return "The user did not allow this change."
                p.parent.mkdir(parents=True, exist_ok=True)
                p.write_text(args["content"], encoding="utf-8")
                return f"Wrote {args['path']}."
            if name == "edit_file":
                p = self.inside(args["path"])
                old = p.read_text(encoding="utf-8")
                if args["old"] not in old:
                    return "The text to replace is not in the file; read it again."
                new = old.replace(args["old"], args["new"], 1)
                self.show_diff(args["path"], old, new)
                if not self.ask(f"edit {args['path']}"):
                    return "The user did not allow this change."
                p.write_text(new, encoding="utf-8")
                return f"Edited {args['path']}."
            if name == "run_command":
                if not self.ask(f"run: {args['command']}"):
                    return "The user did not allow this command."
                done = subprocess.run(args["command"], shell=True, cwd=self.root, capture_output=True, text=True,
                                      timeout=300, encoding="utf-8", errors="replace")
                out = (done.stdout + done.stderr)[-6000:]
                print(dim(textwrap.indent(out.rstrip()[-1500:], "    ")), flush=True)
                return f"exit {done.returncode}\n{out}"
        except subprocess.TimeoutExpired:
            return "The command ran over 5 minutes and was stopped."
        except (OSError, ValueError, KeyError, re.error) as error:
            return f"Failed: {error}"
        return f"Unknown tool {name}."

    @staticmethod
    def show_diff(path: str, old: str, new: str) -> None:
        lines = list(difflib.unified_diff(old.splitlines(), new.splitlines(), path, path, lineterm="", n=1))[:60]
        for line in lines[2:]:
            print(green(line) if line.startswith("+") else red(line) if line.startswith("-") else dim(line), flush=True)


# -- a conversation --------------------------------------------------------------------------------------

AGENT_SYSTEM = ("You are working in the user's project folder through the tools: read before you change, write "
                "whole files, run the build or tests to check, and only say something happened when a tool result "
                "shows it. When the task is done, answer briefly what you did.")
PLAN_SYSTEM = ("Plan mode. Do not carry out the task in this answer. Reply with a plan: one sentence on the goal as "
               "you understand it, then numbered steps -- what will be done in each, then how the result will be "
               "checked. Keep it short. End by asking the user to confirm or change it.")
SUMMARISE = ("Summarise our conversation so far for yourself, so it can continue from the summary alone. Keep: the "
             "user's goals and preferences, decisions made, facts and numbers established, names of files, "
             "functions and commands, code that matters (briefly), and what is still open. Leave out pleasantries. "
             "Write it in the language of the conversation, as compact notes.")


class Session:
    def __init__(self, config: dict, args) -> None:
        self.config = config
        self.key = api_key(config)
        self.level = int(args.level or config.get("level", 5))
        self.temp = int(args.temp or config.get("temperature", 5))
        self.web = bool(args.web or config.get("web", False))
        self.plan = bool(args.plan)
        self.agent = bool(args.agent)
        self.show_think = bool(args.think or config.get("show_think", False))
        self.chips: list[dict] = config.get("chips", [])
        self.chat = Chat()
        self.attachments: list[dict] = []
        self.images: list[dict] = []
        self.hands = Hands(Path.cwd())
        self.last_sent: str | None = None

    # -- one exchange --

    def body(self, messages: list[dict], **extra) -> dict:
        system = [{"role": "system", "content": ch["text"]} for ch in self.chips if ch.get("on")]
        if self.agent:
            system.append({"role": "system", "content": AGENT_SYSTEM + f" The project folder is {self.hands.root}."})
        if self.plan:
            system.append({"role": "system", "content": PLAN_SYSTEM})
        body = {"model": "nimbus-2-apex", "messages": system + messages, "stream": True, "max_tokens": 4096,
                "thinking_level": self.level, "temperature_dial": self.temp, "nimbus_format": True}
        if self.web:
            body["nimbus_tools"] = ["web_search", "fetch_url"]
        if self.agent:
            body["tools"] = TOOLS
        body.update(extra)
        return body

    def stream(self, body: dict, render: bool = True) -> dict:
        """One request, shown as it comes: (answer, reasoning, tool calls, finish, usage, nimbus)."""
        out = {"answer": "", "think": "", "calls": {}, "finish": "", "usage": {}, "nimbus": {}}
        r = Renderer()
        thought_words, think_open, first = 0, False, [False]
        stop_waiting = threading.Event()

        def waiting():
            if not stop_waiting.wait(8):
                if not first[0]:
                    print(dim("  " + w("waking")), flush=True)

        threading.Thread(target=waiting, daemon=True).start()
        try:
            with request("/chat/completions", self.key, body, timeout=900) as response:
                for chunk in events(response):
                    first[0] = True
                    if chunk.get("error"):
                        raise ApiError(str(chunk["error"].get("message") or chunk["error"])[:300])
                    tool = chunk.get("nimbus_tool")
                    if isinstance(tool, dict) and render:
                        if think_open and not self.show_think:
                            print("\r" + " " * 40 + "\r", end="", flush=True)
                        self.show_tool(tool)
                    choice = (chunk.get("choices") or [{}])[0]
                    delta = choice.get("delta") or {}
                    if delta.get("reasoning_content"):
                        piece = delta["reasoning_content"]
                        out["think"] += piece
                        thought_words += len(piece.split())
                        if render:
                            if self.show_think:
                                print(dim(piece), end="", flush=True)
                            else:
                                print("\r" + dim(f"  … {w('thinking')} {thought_words}"), end="", flush=True)
                            think_open = True
                    if delta.get("content"):
                        if think_open and render:
                            print(("\n" if self.show_think else "\r" + " " * 40 + "\r") + dim("  " + w("thought", n=thought_words)), flush=True)
                            think_open = False
                        out["answer"] += delta["content"]
                        if render:
                            r.feed(delta["content"])
                    for call in delta.get("tool_calls") or []:
                        slot = out["calls"].setdefault(call.get("index", 0), {"id": "", "name": "", "arguments": ""})
                        slot["id"] = call.get("id") or slot["id"]
                        fn = call.get("function") or {}
                        slot["name"] = fn.get("name") or slot["name"]
                        slot["arguments"] += fn.get("arguments") or ""
                    if choice.get("finish_reason"):
                        out["finish"] = choice["finish_reason"]
                    if chunk.get("usage"):
                        out["usage"] = chunk["usage"]
                    if chunk.get("nimbus"):
                        out["nimbus"] = chunk["nimbus"]
        finally:
            stop_waiting.set()
            if render:
                if think_open:
                    print(("\n" if self.show_think else "\r" + " " * 40 + "\r") + dim("  " + w("thought", n=thought_words)), flush=True)
                r.end()
        return out

    def show_tool(self, tool: dict) -> None:
        args = tool.get("arguments") or {}
        if tool.get("status") == "running":
            what = (w("searching") + ": " + str(args.get("query", ""))) if tool.get("name") == "web_search" else (w("reading") + ": " + str(args.get("url", "")))
            print(dim("  ⌕ " + what), flush=True)
        elif tool.get("name") == "web_search" and tool.get("results") is not None:
            print(dim("    " + w("results", n=len(tool.get("results") or []))), flush=True)

    def ask(self, text: str, retry: bool = False) -> None:
        if not self.key:
            say(red(w("no_key", page=KEYS_PAGE)))
            return
        if not retry:
            sent = text
            for a in self.attachments:
                sent += f"\n\n--- {a['name']} ---\n{a['text']}"
            message = {"role": "user", "content": text, "sent": sent}
            if self.images:
                message["images"] = [i["name"] for i in self.images]
            self.chat.messages.append(message)
            if not self.chat.title:
                self.chat.title = re.sub(r"\s+", " ", text)[:60]
        history = self.chat.history(int(self.config.get("memory", 12)))
        if self.images and history:
            history[-1] = {"role": "user", "content": [{"type": "text", "text": history[-1]["content"]}]
                           + [{"type": "image_url", "image_url": {"url": i["data"]}} for i in self.images]}
        say()
        started = time.time()
        try:
            got = self.stream(self.body(history))
            steps = 0
            while self.agent and got["finish"] == "tool_calls" and got["calls"] and steps < 40:
                steps += 1
                history.append({"role": "assistant", "content": got["answer"] or None,
                                "tool_calls": [{"id": c_["id"], "type": "function", "function": {"name": c_["name"], "arguments": c_["arguments"]}}
                                               for c_ in got["calls"].values()]})
                for call in got["calls"].values():
                    try:
                        args = json.loads(call["arguments"] or "{}")
                    except ValueError:
                        args = {}
                    label = args.get("path") or args.get("command") or args.get("pattern") or ""
                    say(cyan(f"  ▸ {call['name']}") + (" " + dim(str(label)[:120]) if label else ""))
                    result = self.hands.run(call["name"], args if isinstance(args, dict) else {})
                    history.append({"role": "tool", "tool_call_id": call["id"], "content": result[:12000]})
                say()
                got = self.stream(self.body(history))
        except KeyboardInterrupt:
            say(yellow("\n  " + w("stopped")))
            return
        except ApiError as error:
            say(red(f"  {error}"))
            return
        self.attachments, self.images = [], []
        if not got["answer"] and not got["think"]:
            say(red("  " + w("lost")))
            return
        nimbus = got["nimbus"]
        reply = {"role": "assistant", "content": got["answer"].strip(), "think": got["think"].strip()}
        if nimbus.get("images"):
            reply["vision"] = nimbus.get("vision", "apex-imagexo-v1")
            # The scenes the image module read: the next turns send them as text.
            for m in reversed(self.chat.messages):
                if m.get("role") == "user":
                    m["sent"] = (m.get("sent") or m["content"]) + "".join(f"\n\n[image]\n{s}" for s in nimbus["images"])
                    break
        self.chat.messages.append(reply)
        if isinstance(nimbus.get("context"), dict):
            self.chat.context = nimbus["context"]
        self.chat.save()
        tokens = int((got["usage"] or {}).get("completion_tokens") or 0)
        seconds = time.time() - started
        foot = [f"{seconds:.0f}s"] + ([f"{tokens} tokens", f"{tokens / seconds:.1f}/s"] if tokens else [])
        if reply.get("vision"):
            foot.append(w("image_used"))
        say(dim("  " + " · ".join(foot)))
        self.meter(quiet=True)

    def meter(self, quiet: bool = False) -> None:
        ctx = self.chat.context or {}
        used, limit = int(ctx.get("used") or 0), int(ctx.get("limit") or CONTEXT)
        if not used:
            if not quiet:
                say(dim("  " + w("context", used=0, limit=limit, pct=0)))
            return
        pct = round(100 * used / limit)
        if not quiet or pct >= 70:
            say((yellow if pct >= 70 else dim)("  " + w("context", used=f"{used:,}", limit=f"{limit:,}", pct=pct)))
        if pct >= 85 and quiet:
            self.compact()

    def compact(self) -> None:
        if len([m for m in self.chat.messages if not m.get("summary")]) < 2:
            return
        say(dim("  " + w("compacting")))
        history = self.chat.history(1000) + [{"role": "user", "content": SUMMARISE}]
        level, self.level = self.level, 2
        web, self.web = self.web, False
        agent, self.agent = self.agent, False
        try:
            got = self.stream(self.body(history, max_tokens=1500), render=False)
        except ApiError as error:
            say(red(f"  {error}"))
            return
        finally:
            self.level, self.web, self.agent = level, web, agent
        if got["answer"].strip():
            covered = len(self.chat.messages)
            self.chat.messages.append({"role": "assistant", "content": got["answer"].strip(), "summary": True})
            self.chat.context = {"used": int((got["usage"] or {}).get("completion_tokens") or 600) + 600, "limit": CONTEXT}
            self.chat.save()
            say(green("  " + w("compacted", n=covered)))

    # -- what is typed --

    def attach(self, path: str) -> None:
        p = Path(path.strip().strip('"')).expanduser()
        if not p.is_file():
            say(red(f"  {p}: no such file"))
            return
        kind = mimetypes.guess_type(p.name)[0] or ""
        if kind.startswith("image/"):
            data = p.read_bytes()
            if len(data) > IMAGE_BYTES or kind not in ("image/png", "image/jpeg", "image/webp", "image/gif"):
                say(red("  png, jpeg, webp or gif, up to 6 MB"))
                return
            self.images.append({"name": p.name, "data": f"data:{kind};base64," + base64.b64encode(data).decode()})
        else:
            text = p.read_text(encoding="utf-8", errors="replace")[:MAX_FILE]
            self.attachments.append({"name": p.name, "text": text})
        say(dim("  " + w("attached", name=p.name, size=f"{p.stat().st_size / 1024:.0f} KB")))

    def command(self, line: str) -> bool:
        """A /command; False to leave."""
        global LANG
        name, _, rest = line[1:].partition(" ")
        rest = rest.strip()
        if name in ("exit", "quit", "q"):
            return False
        if name in ("help", "h", "?"):
            say(HELP_FA if LANG == "fa" else HELP)
        elif name == "new":
            self.chat = Chat()
            say(dim("  " + w("new_chat")))
        elif name == "chats":
            for i, c_ in enumerate(Chat.listing()[:20], 1):
                say(f"  {blue(str(i)):>4}  {c_['title'][:60]}  " + dim(f"{c_['n']} · {time.strftime('%m-%d %H:%M', time.localtime(c_['at']))}"))
        elif name == "open":
            listing = Chat.listing()
            try:
                chosen = listing[int(rest) - 1]
                self.chat = Chat(json.loads((CHATS / f"{chosen['id']}.json").read_text(encoding="utf-8")))
                for m in self.chat.messages[-6:]:
                    if m.get("summary"):
                        continue
                    say((bold("› ") + m["content"][:300]) if m["role"] == "user" else "")
                    if m["role"] == "assistant":
                        r = Renderer()
                        r.feed(m["content"])
                        r.end()
                self.meter()
            except (ValueError, IndexError, OSError):
                say(red("  /open N  (N from /chats)"))
        elif name == "rename":
            self.chat.title = rest
            self.chat.save()
        elif name == "delete":
            listing = Chat.listing()
            try:
                (CHATS / f"{listing[int(rest) - 1]['id']}.json").unlink()
            except (ValueError, IndexError, OSError):
                say(red("  /delete N  (N from /chats)"))
        elif name in ("level", "think-level"):
            self.level = max(1, min(20, int(rest))) if rest.isdigit() else self.level
            self.config["level"] = self.level
            save_config(self.config)
            say(dim(f"  thinking {self.level}/20"))
        elif name in ("temp", "temperature"):
            self.temp = max(1, min(10, int(rest))) if rest.isdigit() else self.temp
            self.config["temperature"] = self.temp
            save_config(self.config)
            say(dim(f"  temperature {self.temp}/10"))
        elif name == "web":
            self.web = not self.web
            say(dim(f"  web: {w('on') if self.web else w('off')}"))
        elif name == "plan":
            self.plan = not self.plan
            say(dim(f"  plan: {w('on') if self.plan else w('off')}"))
        elif name == "agent":
            self.agent = not self.agent
            self.hands.always = False
            say(yellow("  " + w("agent_on", cwd=self.hands.root)) if self.agent else dim(f"  agent: {w('off')}"))
        elif name == "think":
            self.show_think = not self.show_think
            say(dim(f"  reasoning shown: {w('on') if self.show_think else w('off')}"))
        elif name in ("file", "image", "attach"):
            self.attach(rest)
        elif name == "compact":
            self.compact()
        elif name == "context":
            self.meter()
        elif name == "retry":
            if self.chat.messages and self.chat.messages[-1]["role"] == "assistant":
                self.chat.messages.pop()
            if self.chat.messages:
                self.ask("", retry=True)
        elif name == "save":
            last = next((m for m in reversed(self.chat.messages) if m["role"] == "assistant"), None)
            if last and rest:
                Path(rest).expanduser().write_text(last["content"], encoding="utf-8")
                say(dim(f"  saved {rest}"))
        elif name == "copy":
            last = next((m for m in reversed(self.chat.messages) if m["role"] == "assistant"), None)
            if last:
                tool = (["clip"] if os.name == "nt" else ["pbcopy"] if sys.platform == "darwin" else ["xclip", "-selection", "clipboard"])
                try:
                    subprocess.run(tool, input=last["content"].encode("utf-16" if os.name == "nt" else "utf-8"), check=True)
                    say(dim("  copied"))
                except (OSError, subprocess.CalledProcessError):
                    say(red("  no clipboard tool here; use /save FILE"))
        elif name == "chips":
            self.chips_command(rest)
        elif name == "lang":
            LANG = "fa" if rest.startswith("fa") else "en"
            self.config["lang"] = LANG
            save_config(self.config)
        elif name == "login":
            self.config = do_login(self.config)
            self.key = api_key(self.config)
        elif name == "status":
            say(dim(f"  {NAME} {VERSION} · {API} · thinking {self.level}/20 · temperature {self.temp}/10 · web {w('on') if self.web else w('off')}"
                    f" · plan {w('on') if self.plan else w('off')} · agent {w('on') if self.agent else w('off')}"))
            self.meter()
        else:
            say(red(f"  unknown command /{name} -- /help"))
        return True

    def chips_command(self, rest: str) -> None:
        verb, _, arg = rest.partition(" ")
        if verb == "add" and arg:
            title, _, text = arg.partition(":")
            self.chips.append({"title": title.strip(), "text": (text or title).strip(), "on": True})
        elif verb in ("on", "off", "del") and arg.isdigit() and 0 < int(arg) <= len(self.chips):
            if verb == "del":
                self.chips.pop(int(arg) - 1)
            else:
                self.chips[int(arg) - 1]["on"] = verb == "on"
        for i, ch in enumerate(self.chips, 1):
            say(f"  {i}. {'●' if ch.get('on') else '○'} {bold(ch['title'])} {dim(ch['text'][:70])}")
        if not self.chips:
            say(dim("  /chips add Title: the instruction    /chips on|off|del N"))
        self.config["chips"] = self.chips
        save_config(self.config)


HELP = """
  /new  /chats  /open N  /rename TITLE  /delete N       conversations (kept in ~/.terminal-cli/chats)
  /level 1-20   /temp 1-10                               Apex's thinking and temperature dials
  /web   /plan   /agent   /think                         web search · plan mode · work on files here · show reasoning
  /file PATH   /image PATH                               attach a file or a picture to the next message
  /compact   /context                                    summarise the conversation · how full it is
  /retry   /save FILE   /copy                            ask again · keep or copy the last answer
  /chips add Title: text | on N | off N | del N          custom instructions sent with every message
  /lang fa|en   /login   /status   /exit
  A line ending in \\ continues on the next one. Ctrl+C stops an answer.
"""
HELP_FA = """
  /new  /chats  /open N  /rename TITLE  /delete N       گفتگوها (در ~/.terminal-cli/chats)
  /level 1-20   /temp 1-10                               دو دکمه‌ی ایپکس: عمق فکر و دما
  /web   /plan   /agent   /think                         جست‌وجوی وب · حالت طرح · کار روی فایل‌های همین پوشه · نمایش فکر
  /file PATH   /image PATH                               پیوست فایل یا تصویر به پیام بعدی
  /compact   /context                                    خلاصه‌سازی گفتگو · میزان پر بودن حافظه
  /retry   /save FILE   /copy                            دوباره بپرس · ذخیره یا کپی آخرین جواب
  /chips add عنوان: متن | on N | off N | del N           دستورهای همیشگی همراه هر پیام
  /lang fa|en   /login   /status   /exit
  خطی که با \\ تمام شود در خط بعد ادامه می‌یابد. Ctrl+C جواب را متوقف می‌کند.
"""


def do_login(config: dict) -> dict:
    say(dim(f"  {KEYS_PAGE}"))
    try:
        import getpass

        key = getpass.getpass(w("key_prompt")).strip()
    except (EOFError, KeyboardInterrupt):
        return config
    if not re.fullmatch(r"nim_[A-Za-z0-9_\-]{20,80}", key):
        say(red(w("key_bad")))
        return config
    config["key"] = key
    save_config(config)
    say(green(w("key_saved", path=CONFIG)))
    return config


def read_line(prompt: str) -> str:
    lines = []
    while True:
        line = input(prompt if not lines else dim("… "))
        if line.endswith("\\"):
            lines.append(line[:-1])
            continue
        lines.append(line)
        return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    global LANG
    ap = argparse.ArgumentParser(prog=NAME, description="Terminal CLI: AhoosAI in the terminal (Nimbus 2 Apex). Terminal CLI made by AhoosAI.",
                                 epilog=f"Keys: {KEYS_PAGE}  ·  Home: https://ahoos-ai.site")
    ap.add_argument("prompt", nargs="*", help="ask once and exit (or 'login')")
    ap.add_argument("--level", type=int, help="thinking level 1-20 (default 5)")
    ap.add_argument("--temp", type=int, help="temperature dial 1-10 (default 5)")
    ap.add_argument("--web", action="store_true", help="let the model search the web")
    ap.add_argument("--plan", action="store_true", help="plan mode: a plan, not the work")
    ap.add_argument("--agent", action="store_true", help="let the model read/write files and run commands here (asked first)")
    ap.add_argument("--think", action="store_true", help="show the reasoning as it comes")
    ap.add_argument("--file", action="append", default=[], help="attach a file or an image (repeatable)")
    ap.add_argument("--json", action="store_true", help="one answer as JSON (answer, reasoning, usage)")
    ap.add_argument("--version", action="version", version=f"{TITLE} {VERSION} (Terminal CLI made by AhoosAI)")
    args = ap.parse_args(argv)
    config = load_config()
    LANG = config.get("lang", "en")
    if args.prompt == ["login"]:
        do_login(config)
        return 0
    session = Session(config, args)
    for f in args.file:
        session.attach(f)
    piped = "" if sys.stdin.isatty() else sys.stdin.read()
    if args.prompt or piped:
        if not session.key:
            say(red(w("no_key", page=KEYS_PAGE)))
            return 2
        text = " ".join(args.prompt) or "Here is the input."
        if piped:
            session.attachments.append({"name": "stdin", "text": piped[:MAX_FILE]})
        if args.json:
            message = text + "".join(f"\n\n--- {a['name']} ---\n{a['text']}" for a in session.attachments)
            try:
                got = session.stream(session.body([{"role": "user", "content": message}]), render=False)
            except ApiError as error:
                print(json.dumps({"error": str(error)}), flush=True)
                return 1
            print(json.dumps({"answer": got["answer"], "reasoning": got["think"], "usage": got["usage"]}, ensure_ascii=False))
            return 0
        session.ask(text)
        return 0
    say(bold(blue(TITLE)) + dim(f" {VERSION} · " + w("hello")))
    if not session.key:
        say(yellow("  " + w("no_key", page=KEYS_PAGE)))
    if session.agent:
        say(yellow("  " + w("agent_on", cwd=session.hands.root)))
    while True:
        try:
            marks = "".join(m for m, on in (("web ", session.web), ("plan ", session.plan), ("agent ", session.agent)) if on)
            line = read_line(blue(dim(marks) + "› "))
        except (EOFError, KeyboardInterrupt):
            say()
            break
        if not line.strip():
            continue
        if line.startswith("/"):
            if not session.command(line.strip()):
                break
            continue
        session.ask(line)
    say(dim("  " + w("bye")))
    return 0


if __name__ == "__main__":
    sys.exit(main())