#!/usr/bin/env python3
"""avalosai — the Avalos AI console: sovereign assistant + ensemble-verified coder, in your terminal.

  avalosai                       open the interactive console (talk to it, or code)
  avalosai login <api_key>       save your key   (~/.avalos/config.json)
  avalosai <file> "<task>"       one-shot: fix/edit a file (verified before applying)
  avalosai doctor                check connection + key

Install:  irm https://avalos.ai/install.ps1 | iex   (registers the `avalosai` command).
Env: AVALOS_API_KEY, AVALOS_BASE (default https://avalos.ai)."""
import sys, os, json, time, threading, random, urllib.request, urllib.error, socket
from pathlib import Path

__version__ = "1.1.0"   # bump on every change that gets published

GOLD = "#FFD600"
COBALT = "#6f93ef"
NAVY = "#0d2b5e"
CFG = Path.home() / ".avalos" / "config.json"

# ── rich (polished TUI) with a graceful plain fallback ───────────────────────
try:
    from rich.console import Console
    from rich.panel import Panel
    from rich.markdown import Markdown
    from rich.text import Text
    from rich import box
    _con = Console()
    _RICH = True
except Exception:
    _RICH = False
    _con = None

# ── prompt_toolkit: real line editing + persistent ↑/↓ history (Windows has no readline) ──
# NOTE: PromptSession() probes the terminal at construction, so it's built LAZILY on first prompt
# (inside the real console) — never at import, which would crash under a pipe / non-console launch.
try:
    from prompt_toolkit import PromptSession
    from prompt_toolkit.history import FileHistory
    from prompt_toolkit.formatted_text import FormattedText
    from prompt_toolkit.styles import Style as _PTStyle
    _ptk_style = _PTStyle.from_dict({
        "bolt": "bold " + GOLD, "name": "bold " + COBALT, "bar": "bold " + GOLD,
        "frame": COBALT,                                 # cobalt input-box border
        "bottom-toolbar": "bg:" + GOLD,                  # persistent gold control bar
        "tb": "bold #0d2b5e bg:" + GOLD,                 # navy text on gold
        "tbsep": "#0d2b5e bg:" + GOLD,
        "tbdim": "#2b3d66 bg:" + GOLD,
    })
    _PTK = True
except Exception:
    _PTK = False
    _ptk_style = None
_ptk = None
_CUR_SID = None                                          # current session, shown in the control bar


def _p(msg="", style=None):
    if _RICH:
        _con.print(msg, style=style)
    else:
        print(msg)


# ── config / http ─────────────────────────────────────────────────────────────
def _cfg():
    try:
        return json.loads(CFG.read_text())
    except Exception:
        return {}


def _save(d):
    CFG.parent.mkdir(parents=True, exist_ok=True)
    CFG.write_text(json.dumps(d, indent=2))


def _base():
    return (os.environ.get("AVALOS_BASE") or _cfg().get("base") or "https://avalos.ai").rstrip("/")


def _key():
    return os.environ.get("AVALOS_API_KEY") or _cfg().get("api_key") or ""


_SESSION_WARNED = False


def _session_token():
    """OPERATOR_CLI_2026_09_18. X-Avalos-Key proves PAID, never OPERATOR:
    _is_operator() (gateway.py:1713) admits only an X-Admin-Token HMAC match or a
    Bearer session JWT whose verified email is sovereign. That is deliberate --
    gateway.py:6024, "a subscription buys compute, not reach" -- so the CLI cleared
    the billing gate and was structurally unable to clear the tool gate, and the
    engine correctly reported having no hands. Mint one with `avalos_auth.py login`.
    Preferred over shipping the admin token because this expires."""
    global _SESSION_WARNED
    c = _cfg()
    tok = os.environ.get("AVALOS_SESSION_TOKEN") or c.get("session_token") or ""
    exp = float(c.get("session_expires") or 0)
    if tok and exp and exp < time.time():
        # Rule 12. Degrading silently here is indistinguishable from never having
        # logged in: the engine just says "tool not available" again and nothing
        # reports why. An expired token is a STATE CHANGE the user must hear about,
        # so say it once per process, then degrade rather than send a dead token.
        if not _SESSION_WARNED:
            _SESSION_WARNED = True
            print("  operator session EXPIRED -- run:  python3 avalos_auth.py login",
                  file=sys.stderr)
            print("  continuing on the paid tier; file and shell tools stay unavailable.",
                  file=sys.stderr)
        return ""
    return tok


def _req(path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    hdrs = {"X-Avalos-Key": _key()}
    _tok = _session_token()
    if _tok:
        hdrs["Authorization"] = "Bearer " + _tok
    if data:
        hdrs["Content-Type"] = "application/json"
    r = urllib.request.Request(_base() + path, data=data, headers=hdrs)
    for attempt in range(3):
        if attempt > 0:
            print(f"  retrying {path} (attempt {attempt + 1})... (stderr)", file=sys.stderr)
            time.sleep(2 if attempt == 1 else 6)   # 2s then 6s, per spec; the generated 2*(2**(a-1)) gave 2s then 4s
        try:
            with urllib.request.urlopen(r, timeout=600) as resp:
                return json.loads(resp.read().decode())
        # HTTPError MUST be caught before URLError -- it is a SUBCLASS of it. The generated
        # version had them the other way round, which made this whole branch dead code: a 4xx
        # was retried as if transient, and a 429 was swallowed into "all 3 attempts failed"
        # instead of reaching the caller. Specific before general, always.
        except urllib.error.HTTPError as e:
            if e.code == 429:
                raise                      # rate limit: surfaced unchanged, never retried here
            if 400 <= e.code < 500:
                raise                      # permanent: retrying cannot change the answer
            if e.code >= 500 and attempt < 2:
                continue                   # transient server fault: retry
            raise RuntimeError("all 3 attempts failed for %s: %s" % (path, e)) from e
        except (socket.timeout, TimeoutError, urllib.error.URLError) as e:
            if attempt == 2:
                raise RuntimeError("all 3 attempts failed for %s: %s" % (path, e)) from e
            continue


# ── working spinner (rotating verbs + elapsed time), à la Claude Code ─────────
_VERBS = ["Thinking", "Weaving", "Pondering", "Reasoning", "Composing", "Divining",
          "Reflecting", "Synthesizing", "Considering", "Kindling"]


class _Working:
    def __init__(self, label=None):
        self.label = label
        self.stop = False
        self.t0 = time.time()
        self._tick = threading.Event()      # interruptible sleep -> __exit__ can join promptly

    def __enter__(self):
        if not _RICH:
            sys.stdout.write("  …working "); sys.stdout.flush(); return self
        self.verb = self.label or random.choice(_VERBS)
        self._live_thread = threading.Thread(target=self._spin, daemon=True)
        self._live_thread.start()
        return self

    def _spin(self):
        from rich.live import Live
        shimmer = [GOLD, "#FFF0A6"]                       # gold ↔ light: the thunderbolt "breathes"
        i = 0
        with Live(console=_con, refresh_per_second=8, transient=True) as live:
            while not self.stop:
                el = int(time.time() - self.t0)
                c = shimmer[i % 2]; i += 1
                t = Text.assemble(("⚡ ", "bold " + c), (self.verb + "… ", c),
                                  ("(%ds · esc to interrupt)" % el, "dim"))
                live.update(t)
                self._tick.wait(0.45)

    def __exit__(self, *a):
        # CLI_CHROME_2026_09_14. `stop = True` then sleep(0.12) did NOT wait for the Live block
        # to close: _spin sits inside time.sleep(0.45), so the transient region could be erased
        # up to 0.45 s later -- on top of the prompt the caller had already drawn. That is the
        # overlapped "Thinking... / ⚡ message" line in the screenshot. Join it instead, and wait
        # on an Event rather than sleeping so the join cannot itself stall for a frame.
        self.stop = True
        try:
            self._tick.set()
        except Exception:
            pass
        if _RICH:
            t = getattr(self, "_live_thread", None)
            if t is not None:
                t.join(timeout=1.5)
        else:
            print()




# ── frame geometry ────────────────────────────────────────────────────────────
# CLI_CHROME_2026_09_14. Every rule in this client is now measured the same way and sized by
# the same function, because the three screenshots that prompted this had three different
# widths on one screen and a top border one cell too long.
def _cols(s):
    """Visible COLUMNS, not characters. The wordmark bolt U+26A1 is East-Asian-wide and paints
    two cells; len() counted it as one, which is precisely why the message rule overshot."""
    import unicodedata
    n = 0
    for ch in s:
        if unicodedata.combining(ch):
            continue
        n += 2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1
    return n


def _frame_w():
    """One width for the banner, the message rule and its closing rule. Shared so they cannot
    drift apart again."""
    import shutil as _sh
    try:
        c = _sh.get_terminal_size((80, 24)).columns
    except Exception:
        c = 80
    return max(32, min(c - 2, 110))




# ── presentation ──────────────────────────────────────────────────────────────
def welcome():
    if not _RICH:
        print("\n  Avalos AI — sovereign assistant + coder  (%s)" % _base())
        print("  Talk to me, or edit code:  file.py : add tests")
        print("  /help · /login <key> · /clear · /exit\n"); return
    body = Text()
    body.append("A V A L O S   A I", style="bold " + GOLD)          # letter-spaced wordmark — our identity
    body.append("\n")
    body.append("sovereign · ensemble-verified · self-hosted", style=COBALT)
    body.append("\n")
    # CLI_CHROME_2026_09_14: was a hardcoded 44 inside a panel as wide as the terminal, so
    # the rule stopped a third of the way across and read as a broken underline.
    body.append("─" * max(20, _frame_w() - 8) + "\n", style=NAVY)                        # navy hairline rule — house style
    body.append("cwd    ", style="dim"); body.append(os.getcwd() + "\n")
    body.append("server ", style="dim"); body.append(_base() + "\n")
    body.append("key    ", style="dim"); body.append((_key()[:10] + "…") if _key() else "not set — /login <key>", style=("" if _key() else "yellow"))
    # cobalt DOUBLE frame + gold ⚡ wordmark riding the top border — deliberately NOT a single grey rounded box
    _con.print(Panel(body, box=box.DOUBLE, border_style=COBALT, padding=(1, 3),
                     title="[bold %s]⚡ avalos.ai[/]" % GOLD, title_align="left"))
    hint = Text()
    hint.append("  talk to me, or edit code:  ", style="dim")
    hint.append("app.py : add input validation", style=COBALT)
    hint.append("     ·  /help /login /clear /exit", style="dim")
    _con.print(hint); _con.print()


def render_reply(text):
    if _RICH:
        try:
            _con.print(Markdown(text))
        except Exception:
            _con.print(text)
        _con.print()
    else:
        print("  " + (text or "").replace("\n", "\n  ") + "\n")


# ── actions ───────────────────────────────────────────────────────────────────
def login(key):
    c = _cfg(); c["api_key"] = key.strip(); _save(c)
    _p("  ✓ key saved", style="green")


def doctor():
    # STALE_PAYLOAD_2026_09_18. The installer and the payload it downloads hash
    # independently, so the published avalos.py silently fell 3 days behind the
    # built one: users ran `install.ps1` and got a build with no operator-session
    # support and no retry-ordering fix, with nothing anywhere reporting the drift.
    # Printing the version makes a stale payload VISIBLE instead of invisible --
    # compare this against what the server expects.
    _p("  avalos-cli " + __version__, style="dim")
    _p("  server " + _base() + "  ·  key " + ((_key()[:10] + "…") if _key() else "(none)"), style="dim")
    _p("  operator session " + ("active" if _session_token() else "none — run: avalos_auth.py login"),
       style="dim")
    try:
        _req("/v1/coder", {"task": "", "content": ""})
        _p("  ✓ reachable", style="green")
    except urllib.error.HTTPError as e:
        _p("  ✓ reachable (server: %s)" % e.code, style="green")
    except Exception as e:
        _p("  ✗ cannot reach: %s" % e, style="red")


def _new_session(title=""):
    r = _req("/v1/sessions", {"title": title})
    sid = r.get("session_id")
    c = _cfg(); c["session_id"] = sid; _save(c)
    return sid


def _ensure_session():
    sid = _cfg().get("session_id")
    if sid:
        try:
            _req("/v1/sessions/" + sid)               # still exists on the server?
            return sid, True
        except Exception:
            pass
    return _new_session(), False


def _session_history(sid, n=6):
    try:
        s = _req("/v1/sessions/" + sid)
        return s.get("messages", [])[-n:], len(s.get("messages", []))
    except Exception:
        return [], 0


def list_sessions():
    try:
        rows = _req("/v1/sessions").get("sessions", [])
    except Exception as e:
        _p("  couldn't list: %s" % e, style="red"); return
    if not rows:
        _p("  no sessions yet", style="dim"); return
    cur = _cfg().get("session_id")
    for r in rows[:12]:
        mark = "→ " if r["sid"] == cur else "  "
        _p("  %s%s  %s  [%d msgs]" % (mark, r["sid"], (r.get("title") or "")[:44], r.get("messages", 0)),
           style=(GOLD if r["sid"] == cur else "dim"))


def chat_session(sid, msg):
    try:
        r = _req("/v1/sessions/%s/chat" % sid, {"message": msg})
        return r.get("reply", "(no reply)")
    except Exception as e:
        return "(couldn't reach Avalos AI: %s)" % e


def code(fp, task):
    p = Path(fp)
    if not p.exists():
        _p("  no such file: %s" % fp, style="red"); return 2
    if not _key():
        _p("  no API key — /login <key>", style="yellow"); return 2
    content = p.read_text(encoding="utf-8")
    try:
        with _Working("Coding"):
            r = _req("/v1/coder", {"task": task, "filename": p.name, "content": content, "mode": "sota"})
            rid = r.get("run_id")
            st, s = "running", {}
            while st == "running":
                time.sleep(5)
                s = _req("/v1/coder/status?rid=" + rid)
                st = s.get("status", "running")
    except Exception as e:
        _p("  error: %s" % e, style="red"); return 2
    if st in ("SUCCESS", "DONE") and s.get("content"):
        bak = p.with_suffix(p.suffix + ".bak"); bak.write_text(content, encoding="utf-8")
        p.write_text(s["content"], encoding="utf-8")
        b = s.get("billing") or {}
        _p("  ✓ %s — applied to %s (backup %s)" % (st, p.name, bak.name), style="green")
        if b:
            _p("    tokens %s · plan %s · api %s" % (b.get("total_tokens", 0), b.get("billed_to_plan", 0), b.get("billed_to_api", 0)), style="dim")
        if s.get("verified") and s.get("proof"):
            if _RICH:
                _con.print(Panel(s["proof"], title="⚡ VERIFIED — passed a generated test in the sandbox",
                                 border_style="green", box=box.ROUNDED, padding=(0, 1)))
            else:
                _p("  ⚡ VERIFIED — passed a generated test in the sandbox:\n" + s["proof"])
        elif st in ("SUCCESS", "DONE"):
            _p("    (applied; the behavioral test wasn't conclusively verified — review the diff)", style="dim")
    else:
        _p("  ✗ %s %s" % (st, s.get("error", "")), style="red")
    return 0


def index_dir(path):
    root = Path(path)
    if not root.exists():
        _p("  no such path: %s" % path, style="red"); return
    exts = {".py", ".js", ".ts", ".tsx", ".jsx", ".go", ".java", ".rs", ".c", ".cpp", ".h", ".hpp",
            ".rb", ".php", ".cs", ".kt", ".swift", ".md", ".sh"}
    skip = {"node_modules", ".git", "venv", ".venv", "__pycache__", "dist", "build", ".next", "target"}
    files = []
    for f in root.rglob("*"):
        if any(part in skip for part in f.parts):
            continue
        if f.is_file() and f.suffix in exts:
            try:
                if f.stat().st_size < 200000:
                    files.append({"path": str(f.relative_to(root)), "content": f.read_text(encoding="utf-8", errors="ignore")})
            except Exception:
                pass
        if len(files) >= 800:
            break
    if not files:
        _p("  no code files found under %s" % path, style="yellow"); return
    with _Working("Indexing"):
        try:
            r = _req("/v1/index", {"files": files})
        except Exception as e:
            _p("  index error: %s" % e, style="red"); return
    _p("  ⚡ indexed %d files → %d chunks (server-side · your codebase is now in-context for chat)"
       % (r.get("files", 0), r.get("chunks", 0)), style="green")


def _toolbar():
    # persistent control bar pinned to the bottom of the terminal — gives the user perspective + control
    sess = _CUR_SID or "no session · /login to persist"
    # CLI_CHROME2_2026_09_14. prompt_toolkit pins bottom_toolbar to the LAST line and draws the
    # input directly above it, so the typing line sat exactly one row off the screen edge and
    # could never drift up -- Carlos: "goes all the way down... should be 1/2 inch up". The
    # toolbar region is sized to its CONTENT, so two leading blank rows float the input clear of
    # the bottom while the control bar still reads as pinned.
    return FormattedText([
        ("", "\n\n"),
        ("class:tb", " ⚡ avalos AI "), ("class:tbsep", "│"),
        ("class:tb", " %s " % sess), ("class:tbsep", "│"),
        ("class:tbdim", " ↑↓ history · file.py : task to edit · /help /exit "),
    ])


def _prompt():
    global _ptk
    if _PTK:                                          # ↑/↓ recalls & edits previous prompts (persistent)
        try:
            if _ptk is None:                          # lazy build — only inside a real console
                _h = Path.home() / ".avalos" / "history"; _h.parent.mkdir(parents=True, exist_ok=True)
                _ptk = PromptSession(history=FileHistory(str(_h)))
            # CLI_CHROME_2026_09_14. Two defects, both visible in Carlos's screenshots.
            #  (a) the old rule subtracted len("╭─ ⚡ message ") = 13, but that string is 14
            #      COLUMNS wide because U+26A1 is East-Asian-wide -- so the border overshot by
            #      one cell every single time. _cols() measures columns, not characters.
            #  (b) it drew a top border and a left gutter and called it a box. prompt_toolkit
            #      wraps the input, so a right edge CANNOT be drawn -- the text ran straight
            #      through where the border was implied to be. An honest titled rule replaces
            #      it, closed by a matching rule once the input returns, at the same width as
            #      the banner.
            w = _frame_w()
            head = "⚡ message "
            top = head + "─" * max(4, w - _cols(head))
            msg = FormattedText([("class:frame", top + "\n"), ("class:frame", " ")])
            _out = _ptk.prompt(msg, style=_ptk_style, bottom_toolbar=_toolbar).strip()
            try:
                _cl = "\x1b[2m" + ("─" * w) + "\x1b[0m\n"   # dim: the closing rule is structure
                _con.file.write(_cl) if _RICH else sys.stdout.write(_cl)
            except Exception:
                pass
            return _out
        except Exception:
            pass                                      # not a real console → fall through to rich/plain
    if _RICH:
        return _con.input(Text.assemble(("⚡ ", "bold " + GOLD), ("avalos ", "bold " + COBALT), ("┃ ", "bold " + GOLD))).strip()
    return input("avalosai> ").strip()


def _is_file_task(line):
    if ":" not in line:
        return None
    left, right = line.split(":", 1)
    left, right = left.strip(), right.strip()
    if left and right and " " not in left and ("." in left or os.path.exists(left)):
        return left, right
    return None


def _show_hist(sid):
    hist, total = _session_history(sid)
    if total:
        _p("  ⤾ resumed %s — %d messages (server-side · survived disconnect)" % (sid, total), style=GOLD)
        for m in hist:
            who = "you   " if m["role"] == "user" else "avalos"
            _p("  %s  %s" % (who, (m.get("content") or "")[:96].replace("\n", " ")), style="dim")
        _p("")


def repl():
    global _CUR_SID
    welcome()
    sid = None
    if _key():
        try:
            sid, resumed = _ensure_session()
            _show_hist(sid) if resumed else _p("  · new session %s\n" % sid, style="dim")
        except Exception:
            sid = None
    else:
        _p("  · /login <key> to enable persistent sessions\n", style="dim")
    while True:
        _CUR_SID = sid                                # keep the control bar current
        try:
            line = _prompt()
        except (EOFError, KeyboardInterrupt):
            _p("\n  bye 🌱  (session saved on the server — reconnect anytime)", style="dim"); return 0
        if not line:
            continue
        low = line.lower().lstrip("/")
        if low in ("exit", "quit", "q"):
            _p("  bye 🌱  (session saved — reconnect anytime)", style="dim"); return 0
        if low in ("help", "?"):
            _p(__doc__, style="dim"); continue
        if low == "doctor":
            doctor(); continue
        if low == "sessions":
            list_sessions(); continue
        if low == "new":
            sid = _new_session(); _p("  · new session %s" % sid, style=GOLD); continue
        if low.startswith("index "):
            index_dir(line.split(None, 1)[1]); continue
        if low.startswith("resume "):
            sid = line.split(None, 1)[1].strip()
            c = _cfg(); c["session_id"] = sid; _save(c); _show_hist(sid); continue
        if low == "clear":
            if _RICH:
                _con.clear()
            welcome(); continue
        if low.startswith("login "):
            login(line.split(None, 1)[1])
            if not sid:
                try:
                    sid, _r = _ensure_session()
                except Exception:
                    pass
            continue
        ft = _is_file_task(line)
        if ft:
            code(ft[0], ft[1]); continue
        with _Working():
            if sid:
                reply = chat_session(sid, line)
            else:
                try:
                    _rr = _req("/v1/chat/completions", {"messages": [{"role": "user", "content": line}]})
                    reply = ((_rr.get("choices") or [{}])[0].get("message", {}) or {}).get("content", "") or "(no reply)"
                except Exception as _e:
                    reply = "(error: %s)" % _e
        render_reply(reply)


def main():
    a = sys.argv[1:]
    if not a:
        return repl()
    if a[0] in ("-h", "--help", "help"):
        _p(__doc__); return 0
    if a[0] == "login" and len(a) >= 2:
        login(a[1]); return 0
    if a[0] == "doctor":
        doctor(); return 0
    if a[0] == "index" and len(a) >= 2:
        index_dir(a[1]); return 0
    if len(a) >= 2:
        return code(a[0], " ".join(a[1:]))
    _p(__doc__); return 0


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