#!/usr/bin/env python3
# Bantam agent, tiny, powerful, DIY. Public domain. Created by Luxferre in 2026.

import sys, os, json, time, subprocess, urllib.request
try: import readline
except ImportError: readline = None

HIST = os.path.expanduser("~/.bantam_history")
SDIR = os.path.expanduser("~/.bantam/sessions")
AUTO = os.path.join(SDIR, "autosave.json")
_COL = False

def c(t, *cs): return t if not _COL or not cs else "\033[" + ";".join(map(str, cs)) + "m" + t + "\033[0m"

def cp(t, *cs):  # colored prompt: hide ANSI escapes from readline's width calc (\001/\002 = RL_PROMPT_START/END_IGNORE)
  if readline is None or not _COL or not cs: return c(t, *cs)
  return "\001\033[" + ";".join(map(str, cs)) + "m\002" + t + "\001\033[0m\002"

def col(cfg):
  if os.environ.get("NO_COLOR") or os.environ.get("BANTAM_NO_COLOR"): return False
  m = str(cfg.get("color", "auto")).strip().lower()
  if m == "always": return True
  if m == "never": return False
  try: return sys.stdout.isatty()
  except Exception: return False

def get_cfg(path="model.cfg"):
  d = {"endpoint": "https://opencode.ai/zen/v1", "model": "deepseek-v4-flash-free", "temperature": "0.7", "api_key": "-", "timeout": "60", "shell_timeout": "120", "max_al_iterations": "1000", "stream": "true", "color": "auto"}
  if os.path.exists(path):
    for ln in open(path, encoding="utf-8"):
      ln = ln.strip()
      if ln and not ln.startswith("#") and "=" in ln:
        k, v = ln.split("=", 1); d[k.strip()] = v.strip()
  if not d.get("api_key") and "OPENAI_API_KEY" in os.environ: d["api_key"] = os.environ["OPENAI_API_KEY"]
  return d

def num(cfg, k, d):
  try: return type(d)(cfg.get(k, d))
  except (TypeError, ValueError): return d

def prompt(path="system.txt"):
  if os.path.exists(path): return open(path, encoding="utf-8").read().strip()
  return "You are Bantam, a tiny, powerful AI agent."

def T(name, desc, props): return {"type": "function", "function": {"name": name, "description": desc, "parameters": {"type": "object", "properties": props, "required": list(props)}}}
TOOLS = [T("shell_exec", "Run a shell command, return output and exit code.", {"command": {"type": "string"}}),
         T("run_subagent", "Run a child agent with a prompt.", {"prompt": {"type": "string"}})]

def shell(cmd, t=120):
  try:
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=float(t))
    return f"{(r.stdout + r.stderr).strip()}\n\nexit: {r.returncode}"
  except subprocess.TimeoutExpired as e:
    out, err = e.stdout or "", e.stderr or ""
    if isinstance(out, bytes): out = out.decode("utf-8", "replace")
    if isinstance(err, bytes): err = err.decode("utf-8", "replace")
    return f"{(out + err).strip()}\n\n[shell timeout after {t}s]\nexit: -1"

fib = [1, 1, 2, 3, 5, 8, 13, 21, 34]

def llm(cfg, msgs, tools):
  url = cfg["endpoint"].rstrip("/") + "/chat/completions"
  h = {"Content-Type": "application/json", "User-Agent": "Mozilla/5.0 (compatible; Bantam/1.0)"}
  k = cfg.get("api_key", "").strip()
  if k and k != "-": h["Authorization"] = "Bearer " + k
  st = cfg.get("stream", "true").lower() in ("true", "1", "yes")
  p = {"model": cfg["model"], "temperature": float(cfg.get("temperature", 0.7)), "messages": msgs, "stream": st}
  if tools: p["tools"] = tools
  pend = c("...requesting...", 1, 2)
  for i, dly in enumerate(fib + [0]):
    try:
      if _COL: sys.stdout.write("\r" + pend); sys.stdout.flush()
      else: sys.stdout.write(pend + "\n"); sys.stdout.flush()
      req = urllib.request.Request(url, data=json.dumps(p).encode(), headers=h, method="POST")
      with urllib.request.urlopen(req, timeout=num(cfg, "timeout", 60)) as r:
        if not st:
          msg = json.loads(r.read().decode())["choices"][0]["message"]
          if _COL: sys.stdout.write("\r\033[K"); sys.stdout.flush()
          return msg
        if _COL: sys.stdout.write("\r\033[K"); sys.stdout.flush()
        content, reas, tcs, rh, ch = "", "", {}, False, False
        for ln in r:
          ln = ln.decode("utf-8").strip()
          if not ln.startswith("data:"): continue
          if ln[5:].strip() == "[DONE]": break
          try:
            dl = json.loads(ln[5:].strip())["choices"][0].get("delta", {})
            rc = dl.get("reasoning_content") or dl.get("reasoning")
            if rc:
              if not rh: sys.stdout.write(c("--- reasoning start ---", 36) + "\n"); rh = True
              sys.stdout.write(c(rc, 2)); sys.stdout.flush(); reas += rc
            cc = dl.get("content")
            if cc:
              if rh and not ch: sys.stdout.write("\n" + c("--- reasoning end ---", 36) + "\n\n")
              ch = True; sys.stdout.write(cc); sys.stdout.flush(); content += cc
            for tc in dl.get("tool_calls", []):
              ti = tc.get("index", 0)
              if ti not in tcs: tcs[ti] = {"id": tc.get("id", ""), "type": "function", "function": {"name": "", "arguments": ""}}
              if tc.get("id"): tcs[ti]["id"] = tc["id"]
              fn = tc.get("function")
              if fn:
                if fn.get("name"): tcs[ti]["function"]["name"] += fn["name"]
                if fn.get("arguments"): tcs[ti]["function"]["arguments"] += fn["arguments"]
          except Exception: pass
        if rh and not ch: sys.stdout.write("\n" + c("--- reasoning end ---", 36) + "\n")
        elif ch: sys.stdout.write("\n")
        sys.stdout.flush()
        m = {"role": "assistant", "content": content or None}
        if reas: m["reasoning_content"] = reas
        if tcs: m["tool_calls"] = list(tcs.values())
        return m
    except Exception as e:
      if _COL: sys.stdout.write("\r\033[K"); sys.stdout.flush()
      if i < len(fib): print(c(f"[network error: {e}, retrying in {dly}s...]", 31)); time.sleep(dly)
      else: raise

MAX_DEPTH = 5

def AL(cfg, msgs, sp, depth=0):
  stime, mx, st = num(cfg, "shell_timeout", 120), num(cfg, "max_al_iterations", 1000), cfg.get("stream", "true").lower() in ("true", "1", "yes")
  for _ in range(mx):
    m = llm(cfg, msgs, TOOLS); msgs.append(m)
    if not st:
      reas = m.get("reasoning_content") or m.get("reasoning")
      if reas: print(c("--- reasoning start ---", 36) + "\n" + c(reas, 2) + "\n" + c("--- reasoning end ---", 36) + "\n")
      if m.get("content"): print(m["content"])
    tcs = m.get("tool_calls")
    if not tcs: break
    for tc in tcs:
      fn, astr = tc["function"]["name"], tc["function"].get("arguments", "{}")
      print(c(f"[tool call: {fn}({astr})]", 33))
      try:
        a = json.loads(astr) if astr else {}
        if not isinstance(a, dict): raise ValueError("args must be a JSON object")
        err = ""
      except Exception as e: err, a = str(e), {}
      if err: res, sty = f"[tool error: invalid JSON args for {fn}: {err}. Raw: {astr!r}]", 31
      elif fn == "shell_exec": res, sty = shell(a.get("command", ""), stime), 2
      elif fn == "run_subagent":
        if depth >= MAX_DEPTH:
          res, sty = f"[subagent depth limit ({MAX_DEPTH}) reached, child not spawned]", 31
        else:
          sub = [{"role": "system", "content": sp + "\n\nImportant: this is a child agent"}, {"role": "user", "content": a.get("prompt", "")}]
          res, sty = last(AL(cfg, sub, sp, depth + 1)), 2
      else: res, sty = f"Unknown tool: {fn}", 31
      print(c(f"[tool result: {fn}]", 32) + "\n" + c(res, sty) + "\n")
      msgs.append({"role": "tool", "tool_call_id": tc["id"], "content": res})
  else: msgs.append({"role": "assistant", "content": f"[max AL iterations ({mx}) reached]"})
  return msgs

def last(msgs):
  for m in reversed(msgs):
    if m.get("role") == "assistant" and m.get("content"): return m["content"]
  return ""

def sdir(): os.makedirs(SDIR, exist_ok=True); return SDIR

def summary(msgs):
  for m in msgs:
    if m.get("role") == "user" and isinstance(m.get("content"), str) and m["content"].strip():
      t = " ".join(m["content"].split()); return t[:80] + ("..." if len(t) > 80 else "")
  return "(empty session)"

def save(msgs):
  d, base, i = sdir(), time.strftime("%Y%m%d-%H%M%S"), 1
  sid, path = base, os.path.join(d, base + ".json")
  while os.path.exists(path): i += 1; sid = base + "-" + str(i); path = os.path.join(d, sid + ".json")
  data = {"id": sid, "created": time.strftime("%Y-%m-%d %H:%M:%S"), "summary": summary(msgs), "messages": msgs}
  open(path, "w", encoding="utf-8").write(json.dumps(data, ensure_ascii=False, indent=2))
  return sid, data["summary"]

def sessions():
  out = []
  if os.path.isdir(SDIR):
    for fn in os.listdir(SDIR):
      if not fn.endswith(".json"): continue
      try:
        d = json.load(open(os.path.join(SDIR, fn), encoding="utf-8"))
        out.append((d.get("id", fn[:-5]), d.get("created", ""), d.get("summary", ""), len(d.get("messages", []))))
      except Exception: pass
  return sorted(out, key=lambda x: x[0], reverse=True)

def load(sid):
  entries = []
  if os.path.isdir(SDIR):
    for fn in os.listdir(SDIR):
      if not fn.endswith(".json"): continue
      try: entries.append((json.load(open(os.path.join(SDIR, fn), encoding="utf-8")).get("id", fn[:-5]), os.path.join(SDIR, fn)))
      except Exception: pass
  hit = next((e for e in entries if e[0] == sid), None)
  if not hit:
    pref = [e for e in entries if e[0].startswith(sid)]
    if len(pref) == 1: hit = pref[0]
    elif len(pref) > 1: raise KeyError("ambiguous prefix: " + ", ".join(e[0] for e in pref))
  if not hit: raise KeyError(sid)
  return json.load(open(hit[1], encoding="utf-8"))["messages"]

def autosave(msgs):
  d = {"id": "autosave", "created": time.strftime("%Y-%m-%d %H:%M:%S"), "summary": summary(msgs), "messages": msgs}
  open(os.path.join(sdir(), "autosave.json"), "w", encoding="utf-8").write(json.dumps(d, ensure_ascii=False, indent=2))

def summarize(cfg, msgs):
  conv = []
  for m in msgs:
    if m.get("role") == "system": continue
    ct = m.get("content")
    if not ct and m.get("tool_calls"): ct = json.dumps([{"function": tc["function"]["name"], "arguments": tc["function"]["arguments"]} for tc in m["tool_calls"]], ensure_ascii=False)
    if not ct: continue
    if len(ct) > 4000: ct = ct[:4000] + "...[truncated]"
    conv.append(f"{m.get('role', '?')}: {ct}")
  if not conv: return "", "no conversation to summarize"
  joined = "\n\n".join(conv)
  if len(joined) > 100000: joined = joined[-100000:] + "\n...[earlier parts truncated]"
  sm = [{"role": "system", "content": "You are a conversation summarizer for an AI agent's context window. Summarize concisely but completely, preserving all important facts, decisions, code, errors, and the current task state, so the agent can continue the work without the original messages. Output only the summary."},
        {"role": "user", "content": "Summarize this conversation:\n\n" + joined}]
  cc = dict(cfg); cc["stream"] = "false"
  try: m = llm(cc, sm, [])
  except Exception as e: return "", f"LLM error: {e}"
  s = (m.get("content") or m.get("reasoning_content") or "").strip()
  return (s, None) if s else ("", "LLM returned an empty summary")

def compact(cfg, msgs):
  if not msgs or msgs[0].get("role") != "system": return msgs, "", "session has no system message"
  s, err = summarize(cfg, msgs)
  if err: return msgs, "", err
  return [{"role": "system", "content": msgs[0]["content"]}, {"role": "user", "content": "Summary of the previous conversation:\n" + s + "\n\nPlease continue from here."}], s, None

def main():
  global _COL
  if readline:
    try: readline.read_history_file(HIST)
    except OSError: pass
    try: readline.parse_and_bind('"\\C-j": "\\C-v\\C-j"')  # real LF via quoted-insert
    except Exception: pass
  sp, cfg = prompt(), get_cfg()
  _COL = col(cfg)
  msgs = [{"role": "system", "content": sp}]
  if len(sys.argv) > 1 and sys.argv[1]:
    p = sys.argv[1]
    if not os.path.exists(p): print(c(f"Error: file '{p}' not found.", 31)); sys.exit(1)
    msgs.append({"role": "user", "content": open(p, encoding="utf-8").read().strip()})
    AL(cfg, msgs, sp); autosave(msgs)
    if readline:
      try: readline.write_history_file(HIST)
      except OSError: pass
    sys.exit(0)
  print(c("Bantam Agent ready", 1, 32) + c("  (Ctrl+J = new line)", 2))
  while True:
    try: u = input(cp("> ", 1, 36)).strip()
    except (EOFError, KeyboardInterrupt): print(); break
    if not u: continue
    if u == "/quit": break
    elif u == "/clear": msgs = [{"role": "system", "content": sp}]; autosave(msgs); continue
    elif u == "/save":
      sid, sm = save(msgs); print(c(f"[session saved: {sid}]", 32) + "  " + c(sm, 2)); continue
    elif u == "/list":
      ss = sessions()
      if not ss: print(c("No sessions saved yet.", 33)); continue
      for sid, st, sm, n in ss:
        mk = c("  (autosave)", 33) if sid == "autosave" else ""
        print(c(sid, 32) + mk + c(f"  {st}  [{n} msgs]", 2) + "\n  " + c(sm, 2))
      continue
    elif u.startswith("/load"):
      parts = u.split(None, 1)
      if len(parts) < 2: print(c("Usage: /load <session-id>", 31)); continue
      try:
        msgs = load(parts[1]); autosave(msgs)
        print(c(f"[session loaded: {parts[1]}]", 32) + "  " + c(summary(msgs), 2))
      except KeyError as e: print(c(f"Session not found: {e}", 31))
      continue
    elif u == "/compact":
      if len(msgs) <= 1: print(c("Nothing to compact yet.", 33)); continue
      print(c("[compacting conversation...]", 33))
      nm, sm, err = compact(cfg, msgs)
      if err: print(c(f"[compact failed: {err}]", 31)); continue
      msgs = nm; autosave(msgs)
      print(c(f"[compacted to {len(msgs)} messages]", 32)); print(c("--- summary ---", 33) + "\n" + c(sm, 2))
      continue
    elif u == "/help":
      print(c("Bantam commands:", 1, 36))
      for k, v in [("/quit", "exit"), ("/clear", "reset to system prompt"), ("/save", "save session"), ("/list", "list sessions"), ("/load <id>", "load session"), ("/compact", "compact context"), ("/help", "show help")]:
        print(c(f"  {k:<12}", 1, 32) + v)
      continue
    msgs.append({"role": "user", "content": u})
    AL(cfg, msgs, sp); autosave(msgs)
  autosave(msgs)
  if readline:
    try: readline.write_history_file(HIST)
    except OSError: pass

if __name__ == "__main__": main()