#!/usr/bin/env python3
"""Persistent REPL kernel for the pi-continual extension.

Protocol: one JSON object per stdin line: {"id": str, "code": str}.
Replies on real stdout as MARKER + JSON: {"id", "ok", "output"}.
User code runs with stdout/stderr captured, so prints cannot corrupt the protocol.
"""

import ast
import io
import json
import os
import re
import subprocess
import sys
import traceback

MARKER = "\x00PIC\x00"
REAL_STDOUT = sys.stdout

ROOT = os.environ.get("PI_CONTINUAL_ROOT") or os.getcwd()
HARNESS_DIR = os.path.join(ROOT, ".pi", "harness")
SKILLS_DIR = os.path.join(ROOT, ".pi", "skills")
SESSION_FILE = os.environ.get("PI_CONTINUAL_SESSION") or ""
PI_BIN = os.environ.get("PI_CONTINUAL_PI_BIN", "pi")


def _send(obj):
    REAL_STDOUT.write(MARKER + json.dumps(obj) + "\n")
    REAL_STDOUT.flush()


def _slug(name):
    slug = re.sub(r"[^a-zA-Z0-9._-]+", "-", str(name).strip()).strip("-").lower()
    if not slug:
        raise ValueError(f"cannot derive a slug from name {name!r}")
    return slug


class RlmHandle:
    """A persistent sub-agent: a full `pi` session that survives between calls."""

    def __init__(self, name):
        self.name = name
        self.dir = os.path.join(HARNESS_DIR, "subagents", _slug(name))
        self.session_dir = os.path.join(self.dir, "sessions")
        os.makedirs(self.session_dir, exist_ok=True)
        self.out_path = os.path.join(self.dir, "output.txt")
        self.proc = None
        self._out = None

    def _has_session(self):
        try:
            return any(f.endswith(".jsonl") for f in os.listdir(self.session_dir))
        except OSError:
            return False

    def send(self, task, agent=None, model=None):
        """Start a new turn in this sub-agent's session. Returns immediately."""
        if self.running():
            raise RuntimeError(
                f"subagent {self.name!r} is still running; call .wait() first"
            )
        args = [PI_BIN, "-p", "--session-dir", self.session_dir]
        if self._has_session():
            args.append("-c")
        if agent:
            spec = os.path.join(HARNESS_DIR, "agents", _slug(agent) + ".md")
            if not os.path.exists(spec):
                raise FileNotFoundError(f"no agent spec {spec!r}; create one with harness.create('agent', ...)")
            args += ["--append-system-prompt", spec]
        if model:
            args += ["--model", model]
        args.append(task)
        self._out = open(self.out_path, "w")
        self.proc = subprocess.Popen(
            args,
            stdout=self._out,
            stderr=subprocess.STDOUT,
            stdin=subprocess.DEVNULL,
            cwd=ROOT,
        )
        return self

    def running(self):
        return self.proc is not None and self.proc.poll() is None

    def result(self):
        """Final output of the last finished turn, or None if still running / never run."""
        if self.proc is None or self.running():
            return None
        try:
            with open(self.out_path) as f:
                return f.read().strip()
        except OSError:
            return None

    def wait(self, timeout=600):
        """Block until the current turn finishes; returns its output."""
        if self.proc is None:
            return None
        self.proc.wait(timeout=timeout)
        return self.result()

    def __repr__(self):
        state = "running" if self.running() else ("idle" if self.proc else "new")
        return f"<subagent {self.name!r} {state}>"


_SUBAGENTS = {}


def rlm(task, name=None, agent=None, model=None):
    """Spawn a sub-agent, or send a follow-up turn to an existing one (same name).

    Non-blocking: returns an RlmHandle immediately. Use .wait() / .result().
    """
    if name is None:
        name = f"agent-{len(_SUBAGENTS) + 1}"
    handle = _SUBAGENTS.get(name)
    if handle is None:
        handle = RlmHandle(name)
        _SUBAGENTS[name] = handle
    return handle.send(task, agent=agent, model=model)


def subagents():
    """All sub-agent handles created in this kernel, by name."""
    return dict(_SUBAGENTS)


class Harness:
    """CRUD over the continual-harness state, stored as markdown files.

    Kinds:
      memory  -> .pi/harness/memory/<name>.md   (injected into the system prompt)
      prompt  -> .pi/harness/prompts/<name>.md  (injected into the system prompt)
      agent   -> .pi/harness/agents/<name>.md   (sub-agent spec, used by rlm(agent=...))
      skill   -> .pi/skills/<name>/SKILL.md     (native pi skill, needs description)
    """

    KIND_DIRS = {
        "memory": os.path.join(HARNESS_DIR, "memory"),
        "prompt": os.path.join(HARNESS_DIR, "prompts"),
        "agent": os.path.join(HARNESS_DIR, "agents"),
    }

    def _path(self, kind, name):
        if kind == "skill":
            return os.path.join(SKILLS_DIR, _slug(name), "SKILL.md")
        if kind not in self.KIND_DIRS:
            raise ValueError(f"unknown kind {kind!r}; use memory/prompt/agent/skill")
        return os.path.join(self.KIND_DIRS[kind], _slug(name) + ".md")

    def create(self, kind, name, content, description=""):
        path = self._path(kind, name)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        if kind == "skill":
            if not description:
                raise ValueError("skills require a description (used for skill discovery)")
            content = f"---\nname: {_slug(name)}\ndescription: {description}\n---\n\n{content}\n"
        with open(path, "w") as f:
            f.write(content if content.endswith("\n") else content + "\n")
        return path

    def update(self, kind, name, content, description=""):
        if not os.path.exists(self._path(kind, name)):
            raise FileNotFoundError(f"no {kind} named {name!r}")
        return self.create(kind, name, content, description)

    def get(self, kind, name):
        with open(self._path(kind, name)) as f:
            return f.read()

    def delete(self, kind, name):
        path = self._path(kind, name)
        os.remove(path)
        if kind == "skill":
            try:
                os.rmdir(os.path.dirname(path))
            except OSError:
                pass
        return path

    def list(self, kind=None):
        if kind is None:
            return {k: self.list(k) for k in ["memory", "prompt", "agent", "skill"]}
        if kind == "skill":
            if not os.path.isdir(SKILLS_DIR):
                return []
            return sorted(
                d for d in os.listdir(SKILLS_DIR)
                if os.path.exists(os.path.join(SKILLS_DIR, d, "SKILL.md"))
            )
        d = self.KIND_DIRS.get(kind)
        if d is None:
            raise ValueError(f"unknown kind {kind!r}")
        if not os.path.isdir(d):
            return []
        return sorted(f[:-3] for f in os.listdir(d) if f.endswith(".md"))


harness = Harness()


def history(n=None):
    """Entries of this session's JSONL (full history, including pre-compaction)."""
    if not SESSION_FILE or not os.path.exists(SESSION_FILE):
        return []
    entries = []
    with open(SESSION_FILE) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    entries.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    return entries[-n:] if n else entries


GLOBALS = {
    "__name__": "__main__",
    "rlm": rlm,
    "subagents": subagents,
    "harness": harness,
    "history": history,
    "ROOT": ROOT,
}


def run(code):
    buf = io.StringIO()
    sys.stdout = buf
    sys.stderr = buf
    try:
        tree = ast.parse(code, mode="exec")
        trailing_expr = None
        if tree.body and isinstance(tree.body[-1], ast.Expr):
            trailing_expr = ast.Expression(tree.body.pop(-1).value)
        exec(compile(tree, "<repl>", "exec"), GLOBALS)
        if trailing_expr is not None:
            value = eval(compile(trailing_expr, "<repl>", "eval"), GLOBALS)
            if value is not None:
                print(repr(value), file=buf)
        ok = True
    except BaseException:
        traceback.print_exc(file=buf)
        ok = False
    finally:
        sys.stdout = REAL_STDOUT
        sys.stderr = sys.__stderr__
    return ok, buf.getvalue()


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            ok, output = run(req.get("code", ""))
            _send({"id": req.get("id"), "ok": ok, "output": output})
        except Exception:
            _send({"id": None, "ok": False, "output": traceback.format_exc()})


if __name__ == "__main__":
    main()