#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
PULSE Gemini Daemon
Monitors Gemini CLI session files in ~/.gemini/tmp/.../chats/
Captures user prompts, model responses, and thoughts in real-time.
"""
import json
import sys
import re
import signal
import time
import os
import glob
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_gemini.log"
TIMESTAMP_FILE = STATE_DIR / "pulse_gemini_last_timestamp.txt"
POLL_INTERVAL = 0.5

# Gemini Path Logic
GEMINI_HOME = Path(os.path.expanduser("~/.gemini"))
GEMINI_TMP = GEMINI_HOME / "tmp"

class PulseGemini:
    def __init__(self, verbose=False):
        self.verbose = verbose
        self.running = True
        self.last_timestamp = ""
        # Session state: {filepath: {"processed_ids": set(), "mtime": float, "size": int}}
        self.sessions = {}
        self.crypto = PulseCrypto()
        self.current_session_file = None
        self.current_mtime = 0.0
        self.current_size = 0
        self.processed_ids = set()
        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 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 _find_latest_session_file(self):
        """Finds the most recently modified session JSON file in Gemini tmp."""
        if not GEMINI_TMP.exists():
            return None
        
        # Look for all .json files in tmp/*/chats/
        pattern = str(GEMINI_TMP / "*" / "chats" / "*.json")
        files = glob.glob(pattern)
        
        if not files:
            return None
            
        # Sort by modification time, newest first
        latest_file = max(files, key=os.path.getmtime)
        return Path(latest_file)

    def _clean_text(self, text, max_len=15000):
        if not text:
            return ""
        # Handle list content (Gemini returns lists of blocks sometimes)
        if isinstance(text, list):
            parts = []
            for block in text:
                if isinstance(block, dict):
                    t = block.get("text", "") or block.get("content", "")
                    if t:
                        parts.append(str(t).strip())
                elif isinstance(block, str):
                    parts.append(block.strip())
            text = " ".join(parts)
        if not isinstance(text, str):
            text = str(text)
        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(self, ts_iso, role, text):
        try:
            # Handle milliseconds in ISO string if present
            ts_iso = ts_iso.replace("Z", "+00:00")
            dt = datetime.fromisoformat(ts_iso)
            ts_str = dt.strftime("%Y-%m-%d %H:%M:%S")
        except ValueError:
            ts_str = str(ts_iso)
        
        content_part = f"[{ts_str}] {role}: {self._clean_text(text)}"
        signature = self.crypto.sign_string(content_part, "gemini_daemon")
        return f"{content_part} | SIG: {signature}"

    def _process_file(self, filepath):
        """Reads the session file and extracts new messages."""
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError):
            return # File might be being written to

        messages = data.get("messages", [])
        new_entries = []

        for msg in messages:
            msg_id = msg.get("id")
            timestamp = msg.get("timestamp")
            
            # Skip if already processed or older than last timestamp
            if msg_id in self.processed_ids:
                continue
            if self.last_timestamp and timestamp <= self.last_timestamp:
                continue

            role = msg.get("type", "").upper()
            content = msg.get("content", "")
            
            # Handle tool calls / thoughts if content is empty but those exist
            if not content:
                if msg.get("thoughts"):
                    content = "[Thinking...]"
                elif msg.get("toolCalls"):
                    content = "[Tool Call]"
            
            if content:
                new_entries.append((timestamp, role, content))
                self.processed_ids.add(msg_id)
                
                # Update last timestamp to the newest message seen
                if timestamp > self.last_timestamp:
                    self.last_timestamp = timestamp

        return new_entries

    def _maybe_refresh_session(self, latest_file):
        if latest_file != self.current_session_file:
            if self.verbose:
                print(f"[Gemini Pulse] New session found: {latest_file.name}")
            self.current_session_file = latest_file
            self.current_mtime = 0.0
            self.current_size = 0
            self.processed_ids = set()

    def _file_changed(self, filepath):
        try:
            stat = filepath.stat()
        except OSError:
            return False
        if stat.st_mtime == self.current_mtime and stat.st_size == self.current_size:
            return False
        self.current_mtime = stat.st_mtime
        self.current_size = stat.st_size
        return True

    def run(self):
        print(f"[Gemini Pulse] Watching {GEMINI_TMP}...")
        
        while self.running:
            latest_file = self._find_latest_session_file()
            
            if latest_file:
                self._maybe_refresh_session(latest_file)
                if not self._file_changed(latest_file):
                    time.sleep(POLL_INTERVAL)
                    continue

                new_entries = self._process_file(latest_file)
                
                if new_entries:
                    with open(LOG_FILE, "a", encoding="utf-8") as f:
                        for ts, role, text in new_entries:
                            log_line = self._format_log_entry(ts, role, redact_pii(text))
                            f.write(log_line + "\n")
                            if self.verbose:
                                print(log_line)
                    self._save_state()
            
            time.sleep(POLL_INTERVAL)

if __name__ == "__main__":
    v = "--verbose" in sys.argv
    PulseGemini(verbose=v).run()
