#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
PULSE Codex Daemon
Monitors Codex CLI rollout JSONL files in ~/.codex/sessions/**/rollout-*.jsonl
Captures user prompts and Codex responses in real-time.
"""
import json
import re
import signal
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

# Add Utilities to path
SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent

# When run directly, ensure pulse package is importable
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from pulse.core.pulse_crypto import PulseCrypto
from pulse.core.pii_filter import redact_pii

# --- Configuration ---
STATE_DIR = ROOT_DIR / "state"

LOG_FILE = STATE_DIR / "pulse_captured_codex.log"
TIMESTAMP_FILE = STATE_DIR / "pulse_codex_last_timestamp.txt"
POLL_INTERVAL = 0.5

CODEX_HOME = Path.home() / ".codex"
SESSIONS_DIR = CODEX_HOME / "sessions"


def extract_text_from_content(content):
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict):
                t = block.get("text", "")
                if t:
                    parts.append(t.strip())
            elif isinstance(block, str):
                parts.append(block.strip())
        return "\n".join([p for p in parts if p])
    return ""


def clean_text(text, max_len=5000):
    text = text.encode("ascii", errors="ignore").decode("ascii")
    text = re.sub(r"\s+", " ", text).strip()
    return text[:max_len] + "..." if len(text) > max_len else text


def format_log_entry(ts_iso, role, text):
    try:
        if isinstance(ts_iso, str):
            dt = datetime.fromisoformat(ts_iso.replace("Z", "+00:00"))
        elif isinstance(ts_iso, (int, float)):
            dt = datetime.fromtimestamp(ts_iso / 1000, tz=timezone.utc)
        else:
            dt = datetime.now(tz=timezone.utc)
        ts_str = dt.strftime("%Y-%m-%d %H:%M:%S")
    except (ValueError, OSError):
        ts_str = str(ts_iso)
    return f"[{ts_str}] {role}: {clean_text(text)}"


class PulseCodex:
    def __init__(self, verbose=False, reprocess=False):
        self.verbose = verbose
        self.reprocess = reprocess
        self.running = True
        self.last_timestamp = ""
        self.file_positions = {}
        self.known_files = set()
        self.crypto = PulseCrypto()
        self._load_state()
        signal.signal(signal.SIGINT, self._handle_signal)
        signal.signal(signal.SIGTERM, self._handle_signal)

    def _handle_signal(self, signum, frame):
        self.running = False

    def _load_state(self):
        if not self.reprocess and TIMESTAMP_FILE.exists():
            text = TIMESTAMP_FILE.read_text(encoding="utf-8").strip()
            if text:
                self.last_timestamp = text

    def _save_state(self):
        if self.last_timestamp:
            TIMESTAMP_FILE.write_text(self.last_timestamp + "\n", encoding="utf-8")

    def _discover_session_files(self):
        if not SESSIONS_DIR.exists():
            return set()
        return set(SESSIONS_DIR.rglob("rollout-*.jsonl"))

    def _read_new_entries(self, filepath):
        entries = []
        try:
            file_size = filepath.stat().st_size
        except OSError:
            return entries
        last_pos = self.file_positions.get(str(filepath), 0)
        if file_size < last_pos:
            last_pos = 0
        if file_size == last_pos:
            return entries

        try:
            with open(filepath, "r", encoding="utf-8") as f:
                f.seek(last_pos)
                new_content = f.read()
                self.file_positions[str(filepath)] = f.tell()
            for line in new_content.splitlines():
                line = line.strip()
                if not line:
                    continue
                try:
                    data = json.loads(line)
                except json.JSONDecodeError:
                    continue

                if data.get("type") != "response_item":
                    continue
                payload = data.get("payload", {})
                if payload.get("type") != "message":
                    continue
                role = payload.get("role")
                if role not in ("user", "assistant"):
                    continue

                timestamp = data.get("timestamp", "")
                if self.last_timestamp and timestamp and timestamp <= self.last_timestamp:
                    continue

                content = extract_text_from_content(payload.get("content", []))
                if not content:
                    continue
                out_role = "USER" if role == "user" else "CODEX"
                entries.append((timestamp, out_role, content))
        except (OSError, UnicodeDecodeError) as e:
            if self.verbose:
                print(f"[Codex Pulse] Warning: {filepath.name}: {e}")
        return entries

    def _append_to_log(self, entries):
        with open(LOG_FILE, "a", encoding="utf-8") as f:
            for ts, role, text in entries:
                content_part = format_log_entry(ts, role, redact_pii(text))
                signature = self.crypto.sign_string(content_part, "codex_daemon")
                f.write(f"{content_part} | SIG: {signature}\n")

    def _initialize_positions(self):
        if self.reprocess or self.last_timestamp:
            return
        files = self._discover_session_files()
        for f in files:
            try:
                self.file_positions[str(f)] = f.stat().st_size
            except OSError:
                pass
        self.known_files = files

    def check_for_new_entries(self):
        all_new = []
        current_files = self._discover_session_files()
        new_files = current_files - self.known_files
        if new_files and self.verbose:
            for nf in new_files:
                print(f"[Codex Pulse] New session file detected: {nf.name}")
        self.known_files = current_files
        for fp in current_files:
            all_new.extend(self._read_new_entries(fp))
        all_new.sort(key=lambda x: x[0] if x[0] else "")
        return all_new

    def run(self):
        print("=" * 55)
        print("  PULSE - Codex Watcher")
        print("=" * 55)
        print(f"  Log:     {LOG_FILE}")
        print(f"  Poll:    every {POLL_INTERVAL}s")
        if self.last_timestamp:
            print(f"  Resume:  {self.last_timestamp}")
        print("\n  Daemon started. Monitoring for new interactions...")
        print("  Press Ctrl+C to stop.\n")
        self._initialize_positions()

        while self.running:
            try:
                new_entries = self.check_for_new_entries()
                if new_entries:
                    self._append_to_log(new_entries)
                    for ts, role, text in new_entries:
                        if self.verbose:
                            print(f"  {role}: {clean_text(text, 80)}")
                        if isinstance(ts, str) and ts > self.last_timestamp:
                            self.last_timestamp = ts
                    self._save_state()
                time.sleep(POLL_INTERVAL)
            except Exception as e:
                if self.verbose:
                    print(f"[Codex Pulse] Error: {e}")
                time.sleep(POLL_INTERVAL)
        self._save_state()


if __name__ == "__main__":
    v = "--verbose" in sys.argv or "-v" in sys.argv
    PulseCodex(verbose=v, reprocess="--reprocess" in sys.argv).run()
