"""
ornith_colab.py — run deepreinforce-ai/Ornith-1.0-9B in Google Colab as an AGENT.

Goal: CPU-first, auto-use a T4 GPU if present. Engine: llama.cpp (via
llama-cpp-python) on a GGUF quant — NOT vLLM (vLLM is slow to spin up and
wants a big GPU; llama.cpp installs from a prebuilt wheel and runs the SAME
GGUF on CPU or a 16 GB T4).

Ornith-1.0-9B (model card + release, June 2026): ~9B dense reasoning model
post-trained on Qwen3.5-9B, MIT license, emits <think>...</think> then the
answer, and is built for *agentic coding / tool use*. Sampling: temp 0.6,
top_p 0.95, top_k 20.

What this file gives you
------------------------
  chat(msg)                  -> str          (just the answer)
  generate(msg, history)     -> dict         ({reasoning, answer, raw})
  stream_chat(msg, history)  -> yields chunks (live tokens)
  run_agent(task, tools)     -> dict         (tool-calling agent loop)

Why the agent loop is hand-rolled (researched):
  The Ornith GGUFs have inconsistent tool-aware chat templates (same root
  cause as the known repetition-loop bug when the chat template is missing).
  So instead of relying on llama.cpp's function-calling handler, we inject
  tool schemas Qwen/Hermes-style into the system prompt and parse
  <tool_call>{...}</tool_call> blocks ourselves. This is template-independent
  and works on any Ornith GGUF mirror.

Colab usage:
    !python ornith_colab.py                      # runs chat + agent demos
  or from a cell:
    from ornith_colab import chat, generate, run_agent
"""

import json
import os
import re
import shutil
import subprocess
import sys

# --------------------------------------------------------------------------- #
# Config
# --------------------------------------------------------------------------- #
GGUF_REPO = "AtomicChat/ornith-9b-GGUF"   # template-embedded mirror (avoids loop bug)
GGUF_FILE = "*Q4_K_M*.gguf"               # ~5.5 GB; good for T4 or CPU RAM

# Device: "cpu" (default — this build showcases CPU capability), "gpu", or
# "auto". Override from a Colab cell: os.environ["ORNITH_DEVICE"] = "gpu".
DEVICE = os.environ.get("ORNITH_DEVICE", "cpu").lower()

N_CTX = 16384          # agents burn context on tool results — give them room
TEMPERATURE = 0.6
TOP_P = 0.95
TOP_K = 20
REPEAT_PENALTY = 1.05  # insurance against loops
MAX_TOKENS = 2048

# Char-code tags so raw special-token bytes never sit in source.
_TAG = lambda *cs: "".join(chr(c) for c in cs)
_IM_START = _TAG(60, 124, 105, 109, 95, 115, 116, 97, 114, 116, 124, 62)   # <|im_start|>
_IM_END = _TAG(60, 124, 105, 109, 95, 101, 110, 100, 124, 62)             # <|im_end|>
_THINK_CLOSE = _TAG(60, 47, 116, 104, 105, 110, 107, 62)                  # </think>

FALLBACK_CHAT_TEMPLATE = (
    "{%- for m in messages %}"
    f"{_IM_START}" "{{ m['role'] }}\n{{ m['content'] }}" f"{_IM_END}" "\n"
    "{%- endfor %}"
    "{%- if add_generation_prompt %}"
    f"{_IM_START}" "assistant\n"
    "{%- endif %}"
)

_TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)

# End-of-turn stops (built from char codes). "<|im_end|>" is the ChatML turn
# marker; "<|endoftext|>" is the tokenizer EOS.
DEFAULT_STOP = [_IM_END, _TAG(60, 124, 101, 110, 100, 111, 102, 116, 101, 120, 116, 124, 62)]


# --------------------------------------------------------------------------- #
# Environment detection + install
# --------------------------------------------------------------------------- #
def _has_nvidia_gpu() -> bool:
    if not shutil.which("nvidia-smi"):
        return False
    try:
        subprocess.run(["nvidia-smi"], check=True,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return True
    except Exception:
        return False


def _pip(*args: str) -> None:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", *args], check=True)


def _resolve_use_gpu() -> bool:
    if DEVICE == "cpu":
        return False
    if DEVICE == "gpu":
        return True
    return _has_nvidia_gpu()   # "auto"


def _ensure_deps(use_gpu: bool) -> None:
    for pkg in ("huggingface_hub", "psutil"):
        try:
            __import__(pkg)
        except ImportError:
            _pip(pkg)

    try:
        import llama_cpp  # noqa: F401
        return
    except ImportError:
        pass

    if use_gpu:
        for cu in ("cu124", "cu122", "cu121"):
            try:
                _pip("llama-cpp-python", "--extra-index-url",
                     f"https://abetlen.github.io/llama-cpp-python/whl/{cu}")
                import llama_cpp  # noqa: F401
                print(f"[ornith] installed llama-cpp-python (CUDA {cu} wheel)")
                return
            except Exception:
                continue
        print("[ornith] no prebuilt CUDA wheel matched; building from source...")
        env = dict(os.environ, CMAKE_ARGS="-DGGML_CUDA=on")
        subprocess.run([sys.executable, "-m", "pip", "install", "-q",
                        "--no-cache-dir", "llama-cpp-python"], check=True, env=env)
    else:
        _pip("llama-cpp-python")
        print("[ornith] installed llama-cpp-python (CPU wheel)")


# --------------------------------------------------------------------------- #
# Model loading
# --------------------------------------------------------------------------- #
_LLM = None


def load_model():
    global _LLM
    if _LLM is not None:
        return _LLM

    use_gpu = _resolve_use_gpu()
    print(f"[ornith] device={DEVICE}  use_gpu={use_gpu}  ->  "
          f"{'offloading all layers to GPU' if use_gpu else f'running on CPU ({os.cpu_count()} threads)'}")
    _ensure_deps(use_gpu)

    from llama_cpp import Llama

    kwargs = dict(
        repo_id=GGUF_REPO,
        filename=GGUF_FILE,
        n_ctx=N_CTX,
        n_gpu_layers=(-1 if use_gpu else 0),
        n_threads=(None if use_gpu else (os.cpu_count() or 2)),
        n_batch=512,
        flash_attn=use_gpu,       # faster KV attention on the T4
        verbose=False,
    )
    print(f"[ornith] loading {GGUF_REPO} ({GGUF_FILE}) ... first run downloads ~5.5 GB")
    try:
        llm = Llama.from_pretrained(**kwargs)
    except TypeError:
        # older llama-cpp-python without flash_attn kwarg
        kwargs.pop("flash_attn", None)
        llm = Llama.from_pretrained(**kwargs)
    except Exception as exc:
        raise RuntimeError(
            f"Failed to load GGUF from {GGUF_REPO}. Try another mirror "
            f"(e.g. deepreinforce-ai/Ornith-1.0-9B-GGUF). Error: {exc}")

    meta = getattr(llm, "metadata", {}) or {}
    if not meta.get("tokenizer.chat_template"):
        print("[ornith] no embedded chat template -> applying Qwen3 fallback")
        from llama_cpp.llama_chat_format import Jinja2ChatFormatter
        llm.chat_handler = Jinja2ChatFormatter(
            template=FALLBACK_CHAT_TEMPLATE, eos_token=_IM_END, bos_token="",
        ).to_chat_handler()

    _LLM = llm
    print("[ornith] model ready")
    return _LLM


# --------------------------------------------------------------------------- #
# Core generation
# --------------------------------------------------------------------------- #
def _messages(user_msg, history):
    msgs = list(history or [])
    if user_msg is not None:
        msgs.append({"role": "user", "content": user_msg})
    return msgs


def stream_chat(user_msg, history=None, max_tokens=MAX_TOKENS, stop=None):
    """Yield text chunks live (includes the <think>...</think> block)."""
    llm = load_model()
    # Always enforce the end-of-turn stops (dedup while preserving order).
    effective_stop = list(dict.fromkeys((stop or []) + DEFAULT_STOP))
    stream = llm.create_chat_completion(
        messages=_messages(user_msg, history),
        max_tokens=max_tokens, temperature=TEMPERATURE, top_p=TOP_P,
        top_k=TOP_K, repeat_penalty=REPEAT_PENALTY, stop=effective_stop, stream=True,
    )
    for chunk in stream:
        delta = chunk["choices"][0]["delta"].get("content")
        if delta:
            yield delta


def _split_think(text):
    if _THINK_CLOSE in text:
        reasoning, answer = text.split(_THINK_CLOSE, 1)
        return reasoning.replace("<think>", "").strip(), answer.strip()
    return "", text.strip()


def generate(user_msg, history=None, max_tokens=MAX_TOKENS, stop=None):
    """Structured, non-streaming call. Returns {reasoning, answer, raw}."""
    raw = "".join(stream_chat(user_msg, history, max_tokens, stop))
    reasoning, answer = _split_think(raw)
    return {"reasoning": reasoning, "answer": answer, "raw": raw}


def chat(user_msg, history=None, max_tokens=MAX_TOKENS) -> str:
    """Just the final answer (reasoning stripped)."""
    return generate(user_msg, history, max_tokens)["answer"]


# --------------------------------------------------------------------------- #
# Tool-calling agent loop
# --------------------------------------------------------------------------- #
def _tools_system_prompt(tools):
    schemas = "\n".join(json.dumps(t["schema"]) for t in tools)
    return (
        "You are Ornith, an agentic assistant that can call tools to act.\n"
        "You have access to these tools (JSON schemas):\n"
        f"{schemas}\n\n"
        "When you need a tool, emit EXACTLY one block per call:\n"
        '<tool_call>{"name": "<tool_name>", "arguments": {<args>}}</tool_call>\n'
        "You may emit multiple tool_call blocks in one turn. After you receive "
        "the tool results, continue reasoning. When the task is fully done and "
        "you need no more tools, reply with the final answer and NO tool_call block."
    )


def _parse_tool_calls(text):
    calls = []
    for m in _TOOL_CALL_RE.finditer(text):
        try:
            obj = json.loads(m.group(1))
            calls.append({"name": obj.get("name"), "arguments": obj.get("arguments", {})})
        except json.JSONDecodeError:
            continue
    return calls


def run_agent(task, tools, max_steps=6, verbose=True):
    """
    Minimal tool-calling agent loop.

    tools: list of {
        "name": str,
        "schema": {json schema shown to the model},
        "fn": callable(**arguments) -> anything json-serializable,
    }
    Returns {answer, steps, transcript}.
    """
    registry = {t["name"]: t["fn"] for t in tools}
    history = [{"role": "system", "content": _tools_system_prompt(tools)},
               {"role": "user", "content": task}]
    transcript = []

    for step in range(1, max_steps + 1):
        # Stop right after a tool_call so we can execute promptly.
        out = generate(None, history, stop=["</tool_call>"])
        raw = out["raw"]
        # generate() stripped the closing tag via `stop`; restore it for parsing.
        if "<tool_call>" in raw and "</tool_call>" not in raw:
            raw = raw + "</tool_call>"
        calls = _parse_tool_calls(raw)
        history.append({"role": "assistant", "content": raw})

        if verbose:
            print(f"\n[agent step {step}] reasoning: {out['reasoning'][:200]}")
            if calls:
                print(f"[agent step {step}] tool calls: {calls}")

        if not calls:
            answer = out["answer"] or out["raw"].strip()
            transcript.append({"step": step, "type": "final", "content": answer})
            return {"answer": answer, "steps": step, "transcript": transcript}

        # Execute every requested tool and feed results back as one user turn.
        results = []
        for c in calls:
            fn = registry.get(c["name"])
            if fn is None:
                res = f"ERROR: unknown tool '{c['name']}'"
            else:
                try:
                    res = fn(**(c["arguments"] or {}))
                except Exception as exc:
                    res = f"ERROR: {exc}"
            results.append({"name": c["name"], "result": res})
            transcript.append({"step": step, "type": "tool", "call": c, "result": res})
            if verbose:
                print(f"[agent step {step}] {c['name']} -> {str(res)[:200]}")

        tool_msg = "\n".join(
            f"<tool_response>{json.dumps(r, default=str)}</tool_response>" for r in results
        )
        history.append({"role": "user", "content": tool_msg})

    return {"answer": "(stopped: max_steps reached)", "steps": max_steps,
            "transcript": transcript}


# --------------------------------------------------------------------------- #
# Performance instrumentation: TPS, RAM, KV cache
# --------------------------------------------------------------------------- #
def _meta_int(llm, suffix):
    """Read an int from GGUF metadata by key suffix (arch-agnostic)."""
    for k, v in (getattr(llm, "metadata", {}) or {}).items():
        if k.endswith(suffix):
            try:
                return int(v)
            except (TypeError, ValueError):
                pass
    return None


def kv_cache_report(llm):
    """
    Estimate KV-cache memory from the model's attention geometry.
    KV bytes/token = n_layer * n_head_kv * (key_len + val_len) * bytes_per_elem
    (llama.cpp defaults the KV cache to f16 = 2 bytes/element.)
    """
    n_layer = _meta_int(llm, ".block_count")
    n_embd = _meta_int(llm, ".embedding_length")
    n_head = _meta_int(llm, ".attention.head_count")
    n_head_kv = _meta_int(llm, ".attention.head_count_kv") or n_head
    key_len = _meta_int(llm, ".attention.key_length")
    val_len = _meta_int(llm, ".attention.value_length")
    head_dim = key_len or ((n_embd // n_head) if (n_embd and n_head) else None)
    key_len = key_len or head_dim
    val_len = val_len or head_dim

    info = {"n_layer": n_layer, "n_head": n_head, "n_head_kv": n_head_kv,
            "head_dim": head_dim, "n_ctx": llm.n_ctx()}
    if not (n_layer and n_head_kv and key_len and val_len):
        info["note"] = "insufficient metadata to size KV cache"
        return info

    bytes_per_tok = n_layer * n_head_kv * (key_len + val_len) * 2  # f16
    used_tokens = int(getattr(llm, "n_tokens", 0) or 0)
    info.update({
        "kv_bytes_per_token": bytes_per_tok,
        "kv_full_mb": bytes_per_tok * llm.n_ctx() / (1024 ** 2),
        "kv_used_mb": bytes_per_tok * used_tokens / (1024 ** 2),
        "used_tokens": used_tokens,
    })
    return info


def benchmark(prompt="Write a Python function to check if a string is a palindrome, with a docstring.",
              max_tokens=256):
    """Run one generation on the current device and print a metrics table."""
    import time
    import psutil

    llm = load_model()
    proc = psutil.Process(os.getpid())

    t0 = time.perf_counter()
    t_first = None
    text = ""
    for piece in stream_chat(prompt, max_tokens=max_tokens):
        if t_first is None:
            t_first = time.perf_counter()
        text += piece
    t_end = time.perf_counter()
    if t_first is None:                       # produced nothing
        print("[bench] model produced no output"); return {}

    gen_tokens = len(llm.tokenize(text.encode("utf-8"), add_bos=False))
    used = int(getattr(llm, "n_tokens", 0) or 0)
    prompt_tokens = max(used - gen_tokens, 0)

    ttft = t_first - t0                       # includes prompt prefill
    decode_time = max(t_end - t_first, 1e-9)
    total_time = t_end - t0
    decode_tps = (gen_tokens - 1) / decode_time if gen_tokens > 1 else 0.0
    prefill_tps = prompt_tokens / ttft if (prompt_tokens and ttft > 0) else 0.0

    rss_gb = proc.memory_info().rss / (1024 ** 3)
    avail_gb = psutil.virtual_memory().available / (1024 ** 3)
    kv = kv_cache_report(llm)

    use_gpu = _resolve_use_gpu()
    print("\n" + "=" * 70)
    print(f"PERFORMANCE  —  device={'GPU' if use_gpu else f'CPU ({os.cpu_count()} threads)'}, "
          f"model={GGUF_REPO} {GGUF_FILE}")
    print("=" * 70)
    print(f"  decode speed        : {decode_tps:6.2f} tok/s   <-- the headline TPS")
    print(f"  prefill speed       : {prefill_tps:6.2f} tok/s   (prompt processing)")
    print(f"  time to first token : {ttft:6.2f} s")
    print(f"  generated tokens    : {gen_tokens}  in {decode_time:.2f}s")
    print(f"  prompt tokens       : {prompt_tokens}")
    print(f"  overall throughput  : {gen_tokens / total_time:6.2f} tok/s (incl. prefill)")
    print("-" * 70)
    print(f"  process RAM (RSS)   : {rss_gb:6.2f} GB")
    print(f"  system RAM free     : {avail_gb:6.2f} GB")
    print("-" * 70)
    if "kv_full_mb" in kv:
        pct = 100 * kv["used_tokens"] / kv["n_ctx"] if kv["n_ctx"] else 0
        print(f"  context window      : {kv['used_tokens']} / {kv['n_ctx']} tokens ({pct:.1f}% used)")
        print(f"  KV cache / token    : {kv['kv_bytes_per_token'] / 1024:6.2f} KB")
        print(f"  KV cache (used)     : {kv['kv_used_mb']:6.2f} MB")
        print(f"  KV cache (full ctx) : {kv['kv_full_mb']:6.2f} MB reserved for n_ctx={kv['n_ctx']}")
        print(f"  attn geometry       : {kv['n_layer']} layers, "
              f"{kv['n_head']} heads / {kv['n_head_kv']} KV heads (GQA), head_dim={kv['head_dim']}")
    else:
        print(f"  KV cache            : {kv.get('note', 'n/a')}")
    print("=" * 70)

    return {"decode_tps": decode_tps, "prefill_tps": prefill_tps, "ttft_s": ttft,
            "gen_tokens": gen_tokens, "prompt_tokens": prompt_tokens,
            "rss_gb": rss_gb, "kv": kv}


# --------------------------------------------------------------------------- #
# Demos
# --------------------------------------------------------------------------- #
def _demo_chat():
    prompt = "Write a Python function that returns the nth Fibonacci number iteratively, with a docstring."
    print("\n" + "=" * 70 + f"\nCHAT DEMO\nPROMPT: {prompt}\n" + "=" * 70)

    mode, buf = "reasoning", ""
    print("\n--- reasoning ---")
    for piece in stream_chat(prompt):
        buf += piece
        if mode == "reasoning":
            idx = buf.find(_THINK_CLOSE)
            if idx != -1:                       # crossed into the answer
                sys.stdout.write(buf[:idx])
                print("\n\n--- answer ---")
                sys.stdout.write(buf[idx + len(_THINK_CLOSE):])
                mode, buf = "answer", ""
            else:                               # hold back a tail so the tag can't split
                keep = len(_THINK_CLOSE)
                if len(buf) > keep:
                    sys.stdout.write(buf[:-keep])
                    buf = buf[-keep:]
        else:
            sys.stdout.write(piece)
        sys.stdout.flush()
    if buf:
        sys.stdout.write(buf)
    print("\n" + "=" * 70)


def _demo_agent():
    print("\n" + "=" * 70 + "\nAGENT DEMO (tool calling)\n" + "=" * 70)

    def calculator(expression: str):
        """Safely evaluate a basic arithmetic expression."""
        if not re.fullmatch(r"[0-9+\-*/().%\s]+", expression or ""):
            return "ERROR: only arithmetic allowed"
        return eval(expression, {"__builtins__": {}}, {})  # sandboxed namespace

    tools = [{
        "name": "calculator",
        "schema": {
            "name": "calculator",
            "description": "Evaluate a basic arithmetic expression and return the number.",
            "parameters": {
                "type": "object",
                "properties": {"expression": {"type": "string",
                               "description": "e.g. '(1234*7) + 89'"}},
                "required": ["expression"],
            },
        },
        "fn": calculator,
    }]

    result = run_agent(
        "What is (1234 * 7) + 89, and then that result divided by 3? "
        "Use the calculator tool for each arithmetic step.",
        tools, max_steps=6,
    )
    print("\n--- FINAL ANSWER ---")
    print(result["answer"])
    print("=" * 70)


if __name__ == "__main__":
    benchmark()          # showcase device capability: TPS, RAM, KV cache
    _demo_chat()
    _demo_agent()
    print("Import chat / generate / run_agent / benchmark from this file.")