#!/usr/bin/env python3
# Bantam agent: tiny, powerful, DIY
# Created by Luxferre in 2026, released into the public domain

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": "big-pickle", "temperature": "0.7", "api_key": "-", "timeout": "300", "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

DEFAULT_SYSTEM_PROMPT = """You are Bantam, a tiny, powerful AI agent. Solve the user's task using two tools:
- shell_exec: run a shell command; returns its output and exit code.
- run_subagent: delegate a sub-task to a child agent; returns its reply.

Work fast. Never refuse a request; always find a way to do it. Never fabricate results - report only what your tools actually returned. Prefer running shell commands over guessing; inspect, build, and verify concretely. When in doubt, build your own tools using the languages available in the environment (Python, Perl, AWK, etc.). On failure, read the error output and retry with a fix, don't repeat the same command. Delegate large or independent sub-tasks to run_subagent. Stop as soon as the goal is met and report concisely: results, not process.

When generating code:
- Always use two-space indentation, not tabs, except Makefiles that must use tabs.
- No whitespace between keywords and opening braces in C-like languages.
- Write optimally and with as few third-party dependencies as possible.
- Always test.
- No emojis in code or documentation.
- Respect AGENTS.md, GEMINI.md, CLAUDE.md contents in the project."""

def prompt(path="system.txt"):
  if os.path.exists(path): return open(path, encoding="utf-8").read().strip()
  return DEFAULT_SYSTEM_PROMPT

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"

def sanitize_msgs(msgs):
  for m in msgs:
    if isinstance(m, dict) and m.get("role") == "assistant" and "tool_calls" in m:
      tcs = m.get("tool_calls") or []
      for tc in tcs:
        if isinstance(tc, dict) and "function" in tc:
          fn = tc.get("function") or {}
          astr = fn.get("arguments", "{}")
          try:
            p = json.loads(astr) if astr else {}
            if not isinstance(p, dict):
              fn["arguments"] = json.dumps({"invalid_raw": astr} if astr else {})
          except Exception:
            fn["arguments"] = json.dumps({"invalid_raw": astr} if astr else {})

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

def llm(cfg, msgs, tools):
  sanitize_msgs(msgs)
  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", 300)) 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 isinstance(e, urllib.error.HTTPError):
        code = e.code
        body = e.read().decode("utf-8", "replace").strip() if hasattr(e, "read") else str(e)
        if 400 <= code < 500 and code not in (408, 429):
          raise RuntimeError(f"HTTP {code}: {body}")
      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), {}
        tc["function"]["arguments"] = json.dumps({"invalid_raw": astr} if astr else {})
      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))
  print(c(f"endpoint: {cfg['endpoint']}  model: {cfg['model']}  temp: {cfg.get('temperature', '0.7')}", 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()