#!/usr/bin/env python3
"""
aos-chat: minimal ArchitectOS chat client

- Talks directly to the OpenAI Chat Completions API.
- Maintains a conversation history in memory.
- Logs conversation to ~/.architectos/chats/.
- Watches assistant messages for [AOS-RUN] JSON blocks.
- Executes those commands locally (MVP: plain subprocess).
- Logs full command output to ~/.architectos/logs/commands/.
- Sends [AOS-RESULT] blocks back as user messages so the model
  can see the result and continue.

- Also supports [AOS-WRITE] blocks to write files directly from
  JSON payloads, avoiding shell-quoting issues for code/config.
- Now supports [AOS-READ] blocks to safely read files and return
  bounded content.

Requirements:
    pip install openai>=1.0.0

Environment:
    OPENAI_API_KEY   - your API key (required)
    AOS_MODEL        - optional, default: "gpt-4o-mini"
"""

import json
import os
import re
import subprocess
import sys
import textwrap
import time
from pathlib import Path
from typing import Any, Dict, List

try:
    from openai import OpenAI
except ImportError:
    print(
        "ERROR: This script requires the 'openai' package (pip install openai>=1.0.0)",
        file=sys.stderr,
    )
    sys.exit(1)


# ============================================================
#  Configuration
# ============================================================

# Where AOS stores its own runtime state (logs, chats, RAG DB, etc.)
# Default: ~/.architectos, override with AOS_HOME
AOS_HOME = Path(os.getenv("AOS_HOME", str(Path.home() / ".architectos"))).expanduser()

CHATS_DIR = AOS_HOME / "chats"
CMD_LOG_DIR = AOS_HOME / "logs" / "commands"

for d in (CHATS_DIR, CMD_LOG_DIR):
    d.mkdir(parents=True, exist_ok=True)

MODEL = os.getenv("AOS_MODEL", "gpt-4o-mini")

MAX_STDOUT_CHARS = 8000
MAX_STDERR_CHARS = 4000

RUN_BLOCK_RE = re.compile(
    r"\[AOS-RUN\](.*?)\[/AOS-RUN\]",
    re.DOTALL | re.IGNORECASE,
)

WRITE_BLOCK_RE = re.compile(
    r"\[AOS-WRITE\](.*?)\[/AOS-WRITE\]",
    re.DOTALL | re.IGNORECASE,
)

READ_BLOCK_RE = re.compile(
    r"\[AOS-READ\](.*?)\[/AOS-READ\]",
    re.DOTALL | re.IGNORECASE,
)

client = OpenAI()


SYSTEM_PROMPT = textwrap.dedent(
    """
    You are the ArchitectOS assistant running in a terminal client called `aos-chat`.

    Environment:
      - You talk to a single human user (Lathem).
      - The client can execute shell commands and write/read files on your behalf.
      - There may be a local RAG helper script at:
          /home/lathem/ArchitectOS/tools/aos_rag.py
        which can index and search previous commands, scripts, and docs.

    To run a shell command, you MUST emit a block like:

      [AOS-RUN]
      {
        "command": "pytest -q",
        "workdir": "/home/lathem/ArchitectOS",
        "timeout_sec": 120
      }
      [/AOS-RUN]

    The client will run the command, capture stdout/stderr/exit_code,
    and respond with a block:

      [AOS-RESULT]
      {
        "command": "...",
        "workdir": "...",
        "exit_code": 0,
        "timeout": false,
        "truncated": false,
        "reason": null,
        "stdout": "...",
        "stderr": "...",
        "log_file": "/home/lathem/ArchitectOS/logs/commands/..."
      }
      [/AOS-RESULT]

    Notes about exit codes and timeouts:
      - If "timeout": true and "reason": "timeout_killed", the process
        was killed by the client after exceeding timeout_sec. This often
        means the program is waiting for input (interactive) or hung.
      - In that case, DO NOT assume resource exhaustion; instead:
          * Consider a non-interactive check (e.g., "ruby -c file.rb",
            "pytest -q", "php -l file.php").
          * Or rerun with piped input (e.g., "printf '3\\n' | ruby script.rb").

    To write or update a text file, you MUST emit a block like:

      [AOS-WRITE]
      {
        "path": "/home/lathem/ArchitectOS/some_script.py",
        "content": "import sys\\nprint('hi')\\n",
        "mode": "w",
        "makedirs": false
      }
      [/AOS-WRITE]

    Notes for AOS-WRITE:
      - 'path' may be absolute or relative; prefer explicit absolute paths.
      - 'content' is the full file contents as a single JSON string with
        normal \\n newlines.
      - 'mode' is "w" (overwrite) or "a" (append). Default is "w".
      - 'makedirs' (optional, default false) will create parent directories
        if they do not exist.
      - The client will reply with an [AOS-RESULT] describing bytes_written,
        path, and any errors.

    To read a text file, you MUST emit a block like:

      [AOS-READ]
      {
        "path": "/home/lathem/ArchitectOS/some_script.py",
        "max_bytes": 16384,
        "encoding": "utf-8"
      }
      [/AOS-READ]

    Notes for AOS-READ:
      - 'path' may be absolute or relative; prefer explicit absolute paths.
      - 'max_bytes' (optional, default 16384) limits the amount of data returned.
        If the file is larger, the content will be truncated and 'truncated': true.
      - 'encoding' (optional, default "utf-8") is used to decode the file.
      - The client will reply with an [AOS-RESULT] describing path, exists,
        size, truncated, content, and any errors.
      - Use AOS-READ to inspect existing files (config, logs, code) BEFORE
        deciding to overwrite them with AOS-WRITE.

    You will see [AOS-RESULT] blocks as USER messages and should:
      * Interpret the results.
      * Explain what happened in plain language.
      * Optionally propose follow-up [AOS-RUN], [AOS-WRITE], or [AOS-READ] operations.

    ------------------------------------------------------------
    Local RAG helper (remember & reuse tools and workflows)
    ------------------------------------------------------------

    There may be a local RAG helper script at:
      /home/lathem/ArchitectOS/tools/aos_rag.py

    It can:
      - Index command logs from: /home/lathem/ArchitectOS/logs/commands
      - Index scripts/docs under: /home/lathem/ArchitectOS
      - Search for relevant prior commands, utilities, and files.

    Examples of using it via AOS-RUN:

      # Index past shell commands from command logs
      [AOS-RUN]
      {
        "command": "python3 tools/aos_rag.py index-commands /home/lathem/ArchitectOS/logs/commands",
        "workdir": "/home/lathem/ArchitectOS",
        "timeout_sec": 120
      }
      [/AOS-RUN]

      # Index Python scripts under the project
      [AOS-RUN]
      {
        "command": "python3 tools/aos_rag.py index-files /home/lathem/ArchitectOS --kind script --ext .py",
        "workdir": "/home/lathem/ArchitectOS",
        "timeout_sec": 300
      }
      [/AOS-RUN]

      # Search for existing formatting/cleanup utilities or workflows
      [AOS-RUN]
      {
        "command": "python3 tools/aos_rag.py search \\"python code formatter or cleanup utility\\"",
        "workdir": "/home/lathem/ArchitectOS",
        "timeout_sec": 120
      }
      [/AOS-RUN]

    When you need to know:
      - which cleanup/formatting utilities already exist,
      - what commands were previously used for similar tasks, or
      - what the project's established style or workflow is,
    you SHOULD consider calling the RAG helper via AOS-RUN instead of guessing.

    Use its results to:
      - reuse existing utilities and scripts,
      - follow prior successful commands,
      - align new code with existing style and patterns.

    ------------------------------------------------------------
    Code editing and cleanup policy (OUTCOME-FOCUSED)
    ------------------------------------------------------------

    When you create, refactor, or significantly modify CODE, you are responsible
    for leaving it in a clean, idiomatic, well-formatted state. The user cares
    more about the final outcome than the specific tools you used.

    In particular, when you touch code:

      1) Inspect before overwrite:
         - Use AOS-READ to look at the existing file so you understand its
           structure and style before replacing large sections.

      2) Prefer automatic formatters or existing utilities:
         - FIRST, look for evidence in prior [AOS-RUN] / [AOS-RESULT] blocks
           in this conversation of commands that look like formatters, e.g.:
             * contain names like "black", "ruff", "isort", "autopep8",
               "yapf", "gofmt", "prettier", "formatter", "format-code".

         - If you have not seen such a command yet, you may:
             * Discover tools by listing likely directories, for example:
                 - "ls ~/ArchitectOS/bin"
                 - "ls ~/ArchitectOS/tools"
             * Use RAG (aos_rag.py) to search for relevant prior commands or scripts.
             * Try standard language-specific formatters (when appropriate), for example:
                 - "black path/to/file.py"
                 - "ruff format path/to/file.py"
                 - "gofmt -w file.go"
                 - "prettier --write file.js"
               using AOS-RUN from the repository root.

         - You do NOT need the user to tell you the formatter name; infer it
           from prior commands, RAG results, or the filesystem. Favor outcome
           over process.

      3) Create or bootstrap a formatter if nothing exists:
         - If no formatter appears available and it is reasonable to do so,
           you may:
             * install a well-known formatter via pip into the existing venv, or
             * write a small, project-local formatting/cleanup script using
               AOS-WRITE (for example, a Python script that normalizes imports,
               indentation, and line length for key files).
         - Document briefly what you created and how to invoke it, so you can
           reuse it later in this session (and index it via RAG if useful).

      4) Manual cleanup as last resort:
         - If external tools are not available or fail, rewrite the code you
           touched in a clean, consistent style yourself:
             * consistent indentation
             * clear naming
             * minimal noise / dead code removed
             * obvious bugs fixed where you see them
         - Then run a syntax check via AOS-RUN (e.g. "python -m py_compile file.py")
           to validate the file.

      5) Post-edit validation:
         - After major edits or cleanup of code:
             * Run an appropriate non-interactive check:
                 - "python -m py_compile file.py"
                 - "pytest -q" for test suites
                 - "ruby -c", "php -l", etc. for other languages.
             * Only skip this if the change is trivially safe (e.g. comments only).

    You should do all of this PROACTIVELY when working with code, even if the
    user does not explicitly say "use the cleanup util". Treat code quality and
    formatting as part of your basic responsibility.

    ------------------------------------------------------------
    Rules for tool use
    ------------------------------------------------------------

      - Only emit ONE tool block ([AOS-RUN], [AOS-WRITE], OR [AOS-READ]) per assistant message.
      - Do not wrap tool blocks in markdown code fences.
      - Keep commands as minimal and safe as possible.
      - Prefer [AOS-WRITE] for code and config files instead of shell 'echo' tricks.
      - Use [AOS-READ] to inspect files (including logs) instead of shell cat/grep
        when you only need raw file content.
      - Avoid running interactive programs (those that use input(), gets, readline,
        etc.) without providing input via a pipe. Prefer:
          * syntax checks (e.g. "ruby -c file.rb", "php -l file.php"), or
          * commands with piped input (e.g. "printf '3\\n' | ruby file.rb").
      - Never run obviously destructive commands (e.g. 'rm -rf /').
      - Assume stdout/stderr may be truncated; for large logs, ask the user
        to open the log_file if you need the full content.

    General behavior:
      - Be concise but not cryptic.
      - Explain your reasoning clearly when doing non-trivial changes.
      - When not using tools, behave like a normal helpful assistant.
    """
).strip()


# ============================================================
#  Logging helpers
# ============================================================

def open_chat_log():
    ts = time.strftime("%Y-%m-%d-%H%M%S")
    path = CHATS_DIR / f"{ts}-aos-chat.log"
    return path.open("a", encoding="utf-8", errors="ignore")


def log_chat(log_fh, role: str, text: str):
    prefix = role.upper()
    for line in text.splitlines():
        log_fh.write(f"[{prefix}] {line}\n")
    log_fh.flush()


def slugify_command(cmd: str, max_len: int = 32) -> str:
    slug = re.sub(r"[^a-zA-Z0-9._-]+", "_", cmd).strip("_")
    if len(slug) > max_len:
        slug = slug[:max_len]
    return slug or "command"


def write_command_log(
    cmd: str,
    stdout: str,
    stderr: str,
    exit_code: int,
    timeout_flag: bool,
) -> Path:
    ts = time.strftime("%Y-%m-%d-%H%M%S")
    slug = slugify_command(cmd)
    path = CMD_LOG_DIR / f"{ts}-{slug}.log"
    with path.open("w", encoding="utf-8", errors="ignore") as fh:
        fh.write(f"COMMAND: {cmd}\n")
        fh.write(f"EXIT_CODE: {exit_code}\n")
        fh.write(f"TIMEOUT: {timeout_flag}\n\n")
        fh.write("=== STDOUT ===\n")
        fh.write(stdout)
        fh.write("\n\n=== STDERR ===\n")
        fh.write(stderr)
    return path


# ============================================================
#  File writer for AOS-WRITE
# ============================================================

def write_file(payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    Handle an [AOS-WRITE] payload.

    Expected keys:
      - path (str): where to write
      - content (str): full file contents
      - mode (str, optional): "w" or "a" (default: "w")
      - makedirs (bool, optional): create parent dirs if needed (default: False)
    """
    path_val = payload.get("path")
    content = payload.get("content")
    mode = payload.get("mode", "w")
    makedirs = bool(payload.get("makedirs", False))

    if not isinstance(path_val, str) or not path_val.strip():
        return {
            "error": "Missing or invalid 'path' in AOS-WRITE payload.",
        }

    if not isinstance(content, str):
        return {
            "error": "Missing or invalid 'content' in AOS-WRITE payload.",
            "path": path_val,
        }

    if mode not in ("w", "a"):
        return {
            "error": "Invalid 'mode' in AOS-WRITE payload (expected 'w' or 'a').",
            "path": path_val,
            "mode": mode,
        }

    try:
        path = Path(path_val).expanduser()
    except Exception as e:
        return {
            "error": f"Failed to resolve path: {e}",
            "path": path_val,
        }

    try:
        if makedirs:
            path.parent.mkdir(parents=True, exist_ok=True)

        data = content
        bytes_len = len(data.encode("utf-8"))

        with path.open(mode, encoding="utf-8", errors="ignore") as fh:
            fh.write(data)

        return {
            "path": str(path),
            "mode": mode,
            "bytes_written": bytes_len,
            "makedirs": makedirs,
            "error": None,
        }
    except Exception as e:
        return {
            "path": str(path),
            "mode": mode,
            "makedirs": makedirs,
            "error": f"Exception while writing file: {e}",
        }


# ============================================================
#  File reader for AOS-READ
# ============================================================

def read_file(payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    Handle an [AOS-READ] payload.

    Expected keys:
      - path (str): which file to read (required)
      - max_bytes (int, optional): limit on bytes to read (default: 16384)
      - encoding (str, optional): text encoding (default: "utf-8")
    """
    path_val = payload.get("path")
    max_bytes = payload.get("max_bytes", 16384)
    encoding = payload.get("encoding", "utf-8") or "utf-8"

    if not isinstance(path_val, str) or not path_val.strip():
        return {
            "error": "Missing or invalid 'path' in AOS-READ payload.",
        }

    try:
        path = Path(path_val).expanduser()
    except Exception as e:
        return {
            "error": f"Failed to resolve path: {e}",
            "path": path_val,
        }

    if not path.exists():
        return {
            "path": str(path),
            "exists": False,
            "size": 0,
            "content": "",
            "truncated": False,
            "error": "File does not exist.",
        }

    if not path.is_file():
        return {
            "path": str(path),
            "exists": True,
            "size": 0,
            "content": "",
            "truncated": False,
            "error": "Path is not a regular file.",
        }

    try:
        size = path.stat().st_size
        # Clamp max_bytes to a sane minimum
        try:
            max_bytes_int = int(max_bytes)
        except (TypeError, ValueError):
            max_bytes_int = 16384

        if max_bytes_int <= 0:
            max_bytes_int = 16384

        truncated = size > max_bytes_int

        # Read from the beginning; if truncated, the assistant can ask
        # for a narrower view (e.g., via grep or smaller file segments).
        with path.open("r", encoding=encoding, errors="replace") as fh:
            content = fh.read(max_bytes_int)

        return {
            "path": str(path),
            "exists": True,
            "size": size,
            "content": content,
            "truncated": truncated,
            "encoding": encoding,
            "error": None,
        }
    except Exception as e:
        return {
            "path": str(path),
            "exists": True,
            "size": 0,
            "content": "",
            "truncated": False,
            "encoding": encoding,
            "error": f"Exception while reading file: {e}",
        }


# ============================================================
#  Command runner for AOS-RUN
# ============================================================

def run_shell_command(payload: Dict[str, Any]) -> Dict[str, Any]:
    cmd = payload.get("command")
    if not cmd or not isinstance(cmd, str):
        return {
            "error": "Missing or invalid 'command' in AOS-RUN payload."
        }

    workdir = payload.get("workdir", os.getcwd())
    timeout_sec = payload.get("timeout_sec", 120)

    if not os.path.isdir(workdir):
        return {
            "command": cmd,
            "workdir": workdir,
            "exit_code": -1,
            "timeout": False,
            "truncated": False,
            "reason": "invalid_workdir",
            "error": f"Workdir does not exist: {workdir}",
            "stdout": "",
            "stderr": "",
        }

    proc = subprocess.Popen(
        ["/bin/bash", "-lc", cmd],
        cwd=workdir,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )

    timeout_flag = False
    try:
        stdout, stderr = proc.communicate(timeout=timeout_sec)
    except subprocess.TimeoutExpired:
        timeout_flag = True
        proc.kill()
        stdout, stderr = proc.communicate()

    exit_code = proc.returncode

    log_path = write_command_log(cmd, stdout, stderr, exit_code, timeout_flag)

    truncated = (
        len(stdout) > MAX_STDOUT_CHARS
        or len(stderr) > MAX_STDERR_CHARS
    )

    reason = None
    if timeout_flag:
        reason = "timeout_killed"
    elif exit_code is None:
        reason = "unknown"
    elif exit_code != 0:
        reason = "nonzero_exit"

    return {
        "command": cmd,
        "workdir": workdir,
        "exit_code": exit_code,
        "timeout": timeout_flag,
        "truncated": truncated,
        "reason": reason,
        "stdout": stdout[-MAX_STDOUT_CHARS:],
        "stderr": stderr[-MAX_STDERR_CHARS:],
        "log_file": str(log_path),
    }


# ============================================================
#  OpenAI chat plumbing
# ============================================================

def call_openai(history: List[Dict[str, str]]) -> str:
    try:
        resp = client.chat.completions.create(
            model=MODEL,
            messages=history,
            temperature=0.2,
        )
    except Exception as e:
        # Keep the REPL alive and surface the error as a pseudo-assistant message.
        err_msg = f"ERROR calling OpenAI API: {e}"
        print(f"\n[aos-chat] {err_msg}\n", file=sys.stderr)
        return err_msg

    return resp.choices[0].message.content or ""


def _handle_tool_block(
    block_text: str,
    kind: str,
) -> Dict[str, Any]:
    """
    Internal helper to parse JSON payload and dispatch to either
    AOS-RUN, AOS-WRITE, or AOS-READ. Returns the result dict.
    """
    try:
        payload = json.loads(block_text)
    except json.JSONDecodeError as e:
        return {
            "error": "Failed to parse JSON in tool block.",
            "exception": str(e),
            "tool": kind,
        }

    if kind == "run":
        return run_shell_command(payload)
    elif kind == "write":
        return write_file(payload)
    elif kind == "read":
        return read_file(payload)
    else:
        return {
            "error": f"Unknown tool kind: {kind}",
        }


def handle_assistant_turn(
    history: List[Dict[str, str]],
    assistant_content: str,
    log_fh,
) -> str:
    """
    Process a single assistant message:
      - log it
      - print it
      - if it contains [AOS-WRITE], [AOS-READ], or [AOS-RUN], execute the operation
        and feed back [AOS-RESULT] in an internal loop until there
        are no more tool blocks.
    Returns the final assistant content (after all tool calls are resolved).
    """

    # First assistant response
    log_chat(log_fh, "assistant", assistant_content)
    print("\nai> " + assistant_content.strip() + "\n")

    while True:
        # Prefer handling file writes first if present (mutating ops),
        # then reads, then runs.
        write_match = WRITE_BLOCK_RE.search(assistant_content)
        read_match = READ_BLOCK_RE.search(assistant_content)
        run_match = RUN_BLOCK_RE.search(assistant_content)

        match = None
        tool_kind = None

        if write_match is not None:
            match = write_match
            tool_kind = "write"
        elif read_match is not None:
            match = read_match
            tool_kind = "read"
        elif run_match is not None:
            match = run_match
            tool_kind = "run"

        if match is None:
            # No tool calls, we're done
            return assistant_content

        block_text = match.group(1).strip()

        # Dispatch to appropriate tool
        result = _handle_tool_block(block_text, tool_kind)
        result_json = json.dumps(result, indent=2)
        result_block = f"[AOS-RESULT]\n{result_json}\n[/AOS-RESULT]"

        history.append({"role": "user", "content": result_block})
        log_chat(log_fh, "aos", result_block)

        # Local hint so the human sees what happened
        if tool_kind == "run":
            exit_code = result.get("exit_code")
            reason = result.get("reason")
            timeout = result.get("timeout")
            print(
                f"\n[aos] ran: {result.get('command')} "
                f"(exit {exit_code}, timeout={timeout}, reason={reason})"
            )
            print(f"[aos] full log: {result.get('log_file')}\n")
        elif tool_kind == "write":
            path = result.get("path", "<unknown>")
            err = result.get("error")
            if err:
                print(f"\n[aos] write FAILED for {path}: {err}\n")
            else:
                print(
                    f"\n[aos] wrote {result.get('bytes_written', 0)} bytes "
                    f"to {path} (mode={result.get('mode')})\n"
                )
        elif tool_kind == "read":
            path = result.get("path", "<unknown>")
            err = result.get("error")
            truncated = result.get("truncated")
            size = result.get("size")
            if err:
                print(f"\n[aos] read FAILED for {path}: {err}\n")
            else:
                print(
                    f"\n[aos] read {len(result.get('content', ''))} bytes "
                    f"from {path} (size={size}, truncated={truncated})\n"
                )

        # Get follow-up assistant message that sees the result
        assistant_content = call_openai(history)
        history.append({"role": "assistant", "content": assistant_content})
        log_chat(log_fh, "assistant", assistant_content)
        print("ai> " + assistant_content.strip() + "\n")


# ============================================================
#  Main REPL
# ============================================================

def repl():
    chat_log = open_chat_log()
    history: List[Dict[str, str]] = [
        {"role": "system", "content": SYSTEM_PROMPT}
    ]

    print(f"[aos-chat] Using model: {MODEL}")
    print("[aos-chat] Type '/quit' or '/exit' to leave.\n")

    try:
        while True:
            try:
                user_input = input("you> ")
            except EOFError:
                break

            user_input = user_input.rstrip("\n")

            if user_input.strip() in ("/quit", "/exit"):
                break

            if not user_input.strip():
                continue

            # Basic out-of-band command to show where logs live
            if user_input.strip() == "/where":
                print(f"  AOS_HOME: {AOS_HOME}")
                print(f"  chats:    {CHATS_DIR}")
                print(f"  cmd logs: {CMD_LOG_DIR}")
                continue

            history.append({"role": "user", "content": user_input})
            log_chat(chat_log, "user", user_input)

            # First assistant turn for this user message
            assistant_content = call_openai(history)
            history.append({"role": "assistant", "content": assistant_content})

            # Handle tool calls if any
            handle_assistant_turn(history, assistant_content, chat_log)

    except KeyboardInterrupt:
        print("\n[aos-chat] Interrupted, exiting.")
    finally:
        chat_log.close()


if __name__ == "__main__":
    if "OPENAI_API_KEY" not in os.environ:
        print("ERROR: OPENAI_API_KEY is not set in the environment.", file=sys.stderr)
        sys.exit(1)

    repl()
