#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary.
"""Log Distiller -- offline batch knowledge extractor for agent logs.

Entry point: python -m pulse.daemons.distiller.distiller [--once] [--verbose]
             [--dry-run] [--reprocess] [--log <path>]
"""
import json
import logging
import re
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional, Tuple

log = logging.getLogger("pulse.distiller")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
MAX_PAIRS_PER_RUN = 100
SESSION_GAP = 1800  # 30 minutes in seconds
POLL_INTERVAL = 60  # seconds between runs in daemon mode
STATE_FILE = Path("state/health/distiller_state.json")

# ---------------------------------------------------------------------------
# Log line parser
# ---------------------------------------------------------------------------
LINE_RE = re.compile(
    r"^\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})\]\s+"
    r"(USER|CLAUDE|GEMINI|CODEX|GROK)\s*:\s+"
    r"(.+?)(?:\s*\|\s*SIG:\s*([a-f0-9]+))?\s*$"
)
SYSTEM_NOISE = re.compile(r"^(SYSTEM BOOT:|# AGENTS\.md|# Context from my IDE|<environment_context|<INSTRUCTIONS)")
SKIP_ROLES = {"INFO", "ERROR", "SYSTEM"}


def _parse_ts(ts_str: str) -> Optional[datetime]:
    """Parse log timestamp to UTC-aware datetime."""
    try:
        return datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
    except ValueError:
        return None


def _parse_lines(raw_text: str) -> List[Tuple[datetime, str, str]]:
    """Parse raw log text into (timestamp, role, content) tuples.

    Filters system noise, skips INFO/ERROR/SYSTEM roles,
    and deduplicates consecutive identical lines.
    """
    parsed = []
    last_content = None
    for line in raw_text.splitlines():
        line = line.rstrip()
        if not line:
            continue
        if SYSTEM_NOISE.match(line):
            continue
        m = LINE_RE.match(line)
        if not m:
            continue
        ts_str, role, content = m.group(1), m.group(2), m.group(3)
        if role in SKIP_ROLES:
            continue
        if content == last_content:
            continue
        last_content = content
        ts = _parse_ts(ts_str)
        if ts is None:
            continue
        parsed.append((ts, role, content))
    return parsed


def _pair_turns(
    lines: List[Tuple[datetime, str, str]],
) -> List[Tuple[datetime, str, str]]:
    """Pair USER lines with subsequent agent lines.

    Returns list of (timestamp, user_text, agent_text).
    Agent text is concatenation of all agent lines following the user line,
    up to the next USER line.
    """
    pairs = []
    i = 0
    while i < len(lines):
        ts, role, content = lines[i]
        if role == "USER":
            agent_parts = []
            j = i + 1
            while j < len(lines) and lines[j][1] != "USER":
                agent_parts.append(lines[j][2])
                j += 1
            agent_text = " ".join(agent_parts)
            pairs.append((ts, content, agent_text))
            i = j
        else:
            i += 1
    return pairs

# ---------------------------------------------------------------------------
# State management (watermark / byte offset)
# ---------------------------------------------------------------------------

def _load_state() -> dict:
    """Load distiller state from state/health/distiller_state.json."""
    try:
        from pulse.core.paths import get_state_dir
        state_path = get_state_dir() / "health" / "distiller_state.json"
    except ImportError:
        state_path = STATE_FILE
    if state_path.exists():
        try:
            return json.loads(state_path.read_text(encoding="utf-8", errors="replace"))
        except Exception:
            pass
    return {}


def _save_state(state: dict) -> None:
    """Persist distiller state atomically."""
    try:
        from pulse.core.paths import get_state_dir
        state_path = get_state_dir() / "health" / "distiller_state.json"
    except ImportError:
        state_path = STATE_FILE
    state_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = state_path.with_suffix(".tmp")
    try:
        tmp.write_text(json.dumps(state, indent=2), encoding="utf-8")
        tmp.replace(state_path)
    except Exception as exc:
        log.debug("state save failed: %s", exc)
        try:
            tmp.unlink(missing_ok=True)
        except Exception:
            pass


def _load_offset(state: dict, log_path: Path) -> int:
    key = f"offset_{log_path.name}"
    return int(state.get(key, 0))


def _save_offset(state: dict, log_path: Path, offset: int) -> None:
    state[f"offset_{log_path.name}"] = offset


# ---------------------------------------------------------------------------
# Log discovery
# ---------------------------------------------------------------------------

def _discover_logs(specific_log: Optional[str] = None) -> List[Path]:
    """Find agent logs in the state directory by globbing pulse_captured_*.log."""
    try:
        from pulse.core.paths import get_state_dir
        state_dir = get_state_dir()
    except ImportError:
        state_dir = Path("state")
    if specific_log:
        p = Path(specific_log)
        return [p] if p.exists() else []
    return sorted(state_dir.glob("pulse_captured_*.log"))


# ---------------------------------------------------------------------------
# Build extractors
# ---------------------------------------------------------------------------

def _build_extractors():
    """Instantiate all extractors. Failures are non-fatal."""
    extractors = []
    specs = [
        ("pulse.daemons.distiller.extractors.knowledge_extractor", "KnowledgeExtractor"),
        ("pulse.daemons.distiller.extractors.wants_extractor", "WantsExtractor"),
        ("pulse.daemons.distiller.extractors.fact_extractor", "FactExtractor"),
        ("pulse.daemons.distiller.extractors.decision_extractor", "DecisionExtractor"),
        ("pulse.daemons.distiller.extractors.reflex_extractor", "ReflexExtractor"),
        ("pulse.daemons.distiller.extractors.llm_extractor", "LLMExtractor"),
    ]
    for module_name, class_name in specs:
        try:
            import importlib
            mod = importlib.import_module(module_name)
            cls = getattr(mod, class_name)
            extractors.append(cls())
        except Exception as exc:
            log.warning("Failed to load %s: %s", class_name, exc)
    return extractors


# ---------------------------------------------------------------------------
# Core processing loop
# ---------------------------------------------------------------------------

def _process_log(
    log_path: Path,
    state: dict,
    extractors: list,
    session_extractor,
    dry_run: bool = False,
    verbose: bool = False,
    reprocess: bool = False,
) -> int:
    """Process a single log file. Returns number of pairs processed."""
    if not log_path.exists():
        return 0
    file_size = log_path.stat().st_size
    offset = 0 if reprocess else _load_offset(state, log_path)
    if file_size < offset:
        log.info("Log rotated for %s -- resetting offset", log_path.name)
        offset = 0
    if offset >= file_size:
        return 0
    raw = log_path.read_bytes()[offset:]
    # Fix 2I: drop partial last line to avoid split-line corruption
    last_nl = raw.rfind(b'\n')
    if last_nl >= 0:
        raw = raw[:last_nl + 1]
    if not raw:
        return 0
    text = raw.decode("utf-8", errors="replace")
    lines = _parse_lines(text)
    pairs = _pair_turns(lines)
    if not pairs:
        if not dry_run:
            _save_offset(state, log_path, offset + len(raw))
        return 0
    source_name = log_path.stem.replace('pulse_captured_', "")
    domain = "general"
    session_turns = []
    last_ts: Optional[datetime] = None
    processed = 0
    for ts, user_text, agent_text in pairs:
        if processed >= MAX_PAIRS_PER_RUN:
            break
        # Fix 2G: update domain per pair from wants_parser
        try:
            from pulse.brain.extraction.wants_parser import parse_full as _wparse
            _parsed = _wparse(user_text, agent_text)
            domain = _parsed.get("primary_domain") or "general"
        except Exception:
            pass
        if last_ts is not None:
            gap = (ts - last_ts).total_seconds()
            if gap > SESSION_GAP:
                if session_turns and not dry_run:
                    try:
                        session_extractor.summarize(
                            session_turns,
                            {"source": source_name, "domain": domain},
                        )
                    except Exception as exc:
                        log.debug("session summarize failed: %s", exc)
                session_turns = []
        metadata = {
            "timestamp": ts.isoformat(),
            "source": source_name,
            "domain": domain,
        }
        if verbose:
            log.info(
                "[%s] pair %d: user=%d chars, agent=%d chars",
                source_name, processed + 1, len(user_text), len(agent_text),
            )
        if not dry_run:
            for extractor in extractors:
                try:
                    items = extractor.extract(user_text, agent_text, metadata)
                    extractor.write(items)
                except Exception as exc:
                    log.debug("%s failed: %s", extractor.name, exc)
        session_turns.append((ts, user_text, agent_text))
        last_ts = ts
        processed += 1
    if session_turns and not dry_run:
        try:
            session_extractor.summarize(
                session_turns,
                {"source": source_name, "domain": domain},
            )
        except Exception as exc:
            log.debug("final session summarize failed: %s", exc)
    if not dry_run:
        _save_offset(state, log_path, offset + len(raw))
    return processed


# CLI entry point lives in distiller_cli.py
if __name__ == "__main__":
    from pulse.daemons.distiller.distiller_cli import main
    main()
