#!/usr/bin/env python3
"""Audited harness wrapper: inject compact output, client, session, and diagnostics.

Explicit CLI options win. Preserve stdout, stderr, and exit code while logging
bounded per-call metadata to the state directory or configured runtime fallback.
Usage: trufflepig-agent [trufflepig options] COMMAND [ARGS...]
"""
from __future__ import annotations

import hashlib
import json
import os
import subprocess
import sys
import time
from pathlib import Path

from trufflepig_runtime import client_environment, fallback_root
from trufflepig_audit_response import coverage_summary, parse_lines_response

SLOW_SECONDS = 3.0
RECENT_LIMIT = 8
SESSION_MARKER_MAX_AGE = 12 * 3600
TRUFFLEPIG = os.environ.get("TRUFFLEPIG_BINARY", "trufflepig")
FLAGS_WITH_VALUES = {
    "--root", "--workspace", "--member", "--cache", "-b", "--budget", "-n",
    "--limit", "--history-cache", "--target", "--side", "--ignore-revs-file",
    "--diagnostics", "--session", "--client", "--format",
}
def state_dir() -> Path:
    base = os.environ.get("XDG_STATE_HOME") or os.path.join(os.path.expanduser("~"), ".local", "state")
    return Path(base) / "trufflepig"


def log_dir() -> Path:
    return Path(os.environ.get("TRUFFLEPIG_AGENT_LOG_DIR") or state_dir() / "agent-audit")


def detect_harness() -> str:
    explicit = os.environ.get("TRUFFLEPIG_AGENT_HARNESS")
    if explicit:
        return explicit
    if os.environ.get("GROK_SESSION_ID"):
        return "grok"
    # omp's shell tool exports CLAUDECODE=1 alongside OMPCODE=1, so omp must be
    # tested before Claude or every omp call would be attributed to Claude.
    if os.environ.get("OMPCODE") == "1":
        return "omp"
    if os.environ.get("CLAUDECODE") == "1" or os.environ.get("CLAUDE_CODE_CHILD_SESSION") == "1":
        return "claude"
    if any(os.environ.get(key) for key in ("CODEX_THREAD_ID", "CODEX_SESSION_ID", "CODEX_VERSION")):
        return "codex"
    for key in os.environ:
        upper = key.upper()
        if "KIMI" in upper:
            return "kimi"
        if "MUSE" in upper:
            return "muse"
    # Walk parent processes and look for a known harness executable name.
    pid = os.getppid()
    for _ in range(12):
        try:
            comm = Path(f"/proc/{pid}/comm").read_text().strip()
            stat = Path(f"/proc/{pid}/stat").read_text()
        except OSError:
            break
        lowered = comm.lower()
        if lowered == "grok" or lowered.startswith("grok-"):
            return "grok"
        if lowered == "claude":
            return "claude"
        if lowered == "codex" or lowered.startswith("codex-"):
            return "codex"
        if lowered == "omp":
            return "omp"
        if lowered.startswith("kimi"):
            return "kimi"
        if lowered.startswith("muse"):
            return "muse"
        try:
            pid = int(stat.rsplit(")", 1)[1].split()[1])
        except (IndexError, ValueError):
            break
        if pid <= 1:
            break
    return "unknown"


def cwd_key(cwd: str) -> str:
    """Per-directory marker name shared with hooks/session-start.sh (Kimi, Muse).

    SHA-256 truncated to 12 hex digits; the hook computes the same digest.
    """
    return hashlib.sha256(cwd.encode()).hexdigest()[:12]


def resolve_session(harness: str, cwd: str) -> str:
    explicit = os.environ.get("TRUFFLEPIG_SESSION")
    if explicit:
        return explicit
    keys = {"grok": ("GROK_SESSION_ID",), "omp": ("TRUFFLEPIG_OMP_SESSION",), "claude": ("TRUFFLEPIG_CLAUDE_SESSION",),
            "codex": ("CODEX_THREAD_ID", "CODEX_SESSION_ID"),
            "kimi": ("KIMI_SESSION_ID",), "muse": ("MUSE_SESSION_ID",)}
    for key in (*keys.get(harness, ()), "AGENT_SESSION_ID"):
        if os.environ.get(key):
            return os.environ[key]
    # Kimi and Muse session hooks write one marker per cwd; Claude and omp
    # arrive through their environment variables above and never read one.
    # A marker older than the age limit is a leftover from a harness that no
    # longer runs its SessionEnd hook, so it is ignored rather than reused.
    key = cwd_key(cwd)
    if harness == "omp":
        return f"{harness}-{key}-{time.strftime('%Y%m%d')}"
    marker = state_dir() / "agent-sessions" / harness / key
    try:
        fresh = time.time() - marker.stat().st_mtime < SESSION_MARKER_MAX_AGE
        session = marker.read_text().strip() if fresh else ""
        if session:
            return session
    except OSError:
        pass
    return f"{harness}-{key}-{time.strftime('%Y%m%d')}"


def split_args(argv: list[str]) -> tuple[list[str], str, list[str], dict[str, str]]:
    """Return (options, verb, words, option values) without altering argv."""
    options: list[str] = []
    values: dict[str, str] = {}
    words: list[str] = []
    index = 0
    while index < len(argv):
        arg = argv[index]
        if arg == "--":
            words.extend(argv[index + 1:])
            break
        if arg.startswith("-") and len(arg) > 1:
            if "=" in arg and arg.startswith("--"):
                name, value = arg.split("=", 1)
                values[name] = value
                options.append(arg)
            elif arg in FLAGS_WITH_VALUES and index + 1 < len(argv):
                values[arg] = argv[index + 1]
                options.extend(argv[index:index + 2])
                index += 1
            elif arg.startswith("-b") and not arg.startswith("--"):
                values["-b"] = arg[2:]
                options.append(arg)
            elif arg.startswith("-n") and not arg.startswith("--"):
                values["-n"] = arg[2:]
                options.append(arg)
            else:
                options.append(arg)
            index += 1
            continue
        words.append(arg)
        index += 1
    # Unknown first words are recorded verbatim so usage errors stay visible.
    verb = words[0] if words else "status"
    return options, verb, words, values


def recent_fingerprints(path: Path) -> list[str]:
    try:
        return path.read_text().split()[-RECENT_LIMIT:]
    except OSError:
        return []


def main() -> int:
    argv = sys.argv[1:]
    cwd = os.getcwd()
    harness = detect_harness()
    session = resolve_session(harness, cwd)
    options, verb, words, values = split_args(argv)
    harness = values.get("--client", harness)
    session = values.get("--session", session)
    injected: list[str] = []
    if "--json" not in options and "--format" not in values:
        injected.extend(["--format", "lines"])
    if "--client" not in values:
        injected.extend(["--client", harness])
    if "--session" not in values:
        injected.extend(["--session", session])
    if "--diagnostics" not in values:
        injected.extend(["--diagnostics", os.environ.get("TRUFFLEPIG_AGENT_DIAGNOSTICS", "detailed")])
    command = [TRUFFLEPIG, *injected, *argv]

    started = time.monotonic()
    try:
        completed = subprocess.run(command, capture_output=True, check=False, env=client_environment())
        exit_code = completed.returncode
        stdout, stderr = completed.stdout, completed.stderr
    except FileNotFoundError:
        exit_code = 127
        stdout = b""
        stderr = f"trufflepig-agent: {TRUFFLEPIG} not found on PATH\n".encode()
    elapsed = time.monotonic() - started

    sys.stdout.buffer.write(stdout)
    sys.stdout.buffer.flush()
    sys.stderr.buffer.write(stderr)
    sys.stderr.buffer.flush()

    response: dict = {}
    if stdout.lstrip().startswith(b"{"):
        try:
            first_line = stdout.split(b"\n", 1)[0]
            parsed = json.loads(first_line) if first_line.strip() else {}
            response = parsed if isinstance(parsed, dict) else {}
        except (ValueError, UnicodeDecodeError):
            response = {}
    elif stdout.strip():
        response = parse_lines_response(stdout)

    hits = response.get("hits")
    hit_count = len(hits) if isinstance(hits, list) else None
    status = response.get("status") or ("ok" if exit_code == 0 else "error")
    error_text = ""
    if isinstance(response.get("error"), str):
        error_text = response["error"]
    stderr_head = stderr.decode("utf-8", "replace").strip()[:300]
    coverage = coverage_summary(response)

    fingerprint = hashlib.blake2s(" ".join([verb, *words[1:]]).encode(), digest_size=8).hexdigest()
    recent_path = state_dir() / "agent-recent" / harness / f"{cwd_key(cwd)}-{session}"
    runtime = fallback_root()
    audit_fallback = runtime / "agent-audit"
    recent_fallback = runtime / "agent-recent" / recent_path.parent.name / recent_path.name
    try:
        recent = recent_path.read_text().split()[-RECENT_LIMIT:]
    except OSError:
        recent = recent_fingerprints(recent_fallback)
    repeat_count = recent.count(fingerprint)

    signals: list[str] = []
    if exit_code != 0 or status not in ("ok",):
        signals.append("error")
    if verb == "search" and hit_count == 0 and exit_code == 0:
        signals.append("no_hits")
    if response.get("truncated") is True:
        signals.append("truncated")
    lowered = (error_text + " " + stderr_head).lower()
    if "budget" in lowered:
        signals.append("budget")
    if any(key in lowered for key in ("stale_result", "stale_source", "expired_result")):
        signals.append("stale")
    if "unknown_command" in lowered or "invalid_query" in lowered or "usage" in lowered:
        signals.append("usage_error")
    if any(m.get("semantic_status") in ("unavailable", "partial") for m in coverage.values()):
        signals.append("semantic_degraded")
    if any(m.get("rerank_status") == "unavailable" for m in coverage.values()):
        signals.append("rerank_unavailable")
    if elapsed > SLOW_SECONDS:
        signals.append("slow")
    if repeat_count >= 1 and verb in ("search", "refs", "map"):
        signals.append("repeat_query")

    record = {
        "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "harness": harness,
        "session": session,
        "cwd": cwd,
        "verb": verb,
        "args": words[1:],
        "options": {k: v for k, v in values.items() if k not in ("--session", "--client")},
        "exit_code": exit_code,
        "elapsed_ms": int(elapsed * 1000),
        "status": status,
        "error": error_text[:300],
        "stderr": stderr_head,
        "hits": hit_count,
        "truncated": response.get("truncated"),
        "has_next": bool(response.get("next")),
        "stdout_bytes": len(stdout),
        "coverage": coverage,
        "repeat_count": repeat_count,
        "signals": signals,
    }
    def append_audit(log_target: Path) -> None:
        log_target.mkdir(parents=True, exist_ok=True)
        with (log_target / f"{harness}.jsonl").open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, separators=(",", ":")) + "\n")

    def write_recent(recent_file: Path) -> None:
        recent_file.parent.mkdir(parents=True, exist_ok=True)
        recent_file.write_text(" ".join([*recent, fingerprint][-RECENT_LIMIT:]))

    def with_fallback(write, primary: Path, fallback: Path) -> OSError | None:
        try:
            write(primary)
            return None
        except OSError:
            try:
                write(fallback)
                return None
            except OSError as fallback_error:
                return fallback_error

    audit_error = with_fallback(append_audit, log_dir(), audit_fallback)
    with_fallback(write_recent, recent_path, recent_fallback)
    if audit_error is not None:
        sys.stderr.write(f"trufflepig-agent: audit log unavailable: {audit_error}\n")
    return exit_code


if __name__ == "__main__":
    sys.exit(main())
