#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Backfill Session Descriptions: Fix 1,927 content-less stubs in current_session.json.

Stubs have description="Interaction captured" or missing type/source fields.
This script matches them against raw log files and agent_sessions.json
by timestamp (5s fuzzy window) to recover actual content.

Usage:
    python backfill_session_descriptions.py --dry-run    # preview changes
    python backfill_session_descriptions.py --execute     # apply changes
"""

import argparse
import io
import json
import os
import re
import sys
from datetime import datetime, timezone, timedelta
from pathlib import Path

if os.name == 'nt' and __name__ == '__main__':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')


from pulse.core.brain_utils import save_json, load_json, EVIDENCE_DIR

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent.parent
STATE_DIR = ROOT_DIR / "state"
SESSION_FILE = EVIDENCE_DIR / "Sessions" / "current_session.json"
AGENT_SESSIONS_FILE = EVIDENCE_DIR / "Sessions" / "agent_sessions.json"

# Log line: [YYYY-MM-DD HH:MM:SS] ROLE: text | SIG: hash
_LOG_LINE = re.compile(
    r"^\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})\] (USER|CLAUDE|GEMINI|CODEX|GROK|ASSISTANT): (.+?)(?:\s*\| SIG: [a-f0-9]+)?$"
)

STUB_DESCRIPTIONS = {"Interaction captured", "", None}
FUZZY_WINDOW_S = 5


def _parse_ts(ts_str):
    """Parse an ISO or log timestamp to datetime."""
    if not ts_str:
        return None
    try:
        s = ts_str.replace("Z", "+00:00")
        if "T" in s:
            return datetime.fromisoformat(s)
        return datetime.strptime(s, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
    except (ValueError, TypeError):
        return None


def _build_log_index():
    """Build timestamp->text index from raw log files."""
    index = []  # list of (datetime, user_text, agent_text)
    for log_file in sorted(STATE_DIR.glob("pulse_captured_*.log")):
        try:
            text = log_file.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue

        user_buf = None
        user_ts = None
        agent_lines = []

        for line in text.splitlines():
            line = line.strip()
            if not line:
                continue
            content = line.rsplit(" | SIG: ", 1)[0] if " | SIG: " in line else line
            m = _LOG_LINE.match(content)
            if not m:
                continue
            ts_str, role, content_text = m.groups()
            ts = _parse_ts(ts_str)
            if not ts:
                continue

            if role == "USER":
                if user_buf and agent_lines:
                    index.append((user_ts, user_buf, "\n".join(agent_lines)))
                user_buf = content_text
                user_ts = ts
                agent_lines = []
            elif user_buf is not None:
                agent_lines.append(content_text)

        # Flush last
        if user_buf and agent_lines:
            index.append((user_ts, user_buf, "\n".join(agent_lines)))

    return index


def _build_session_index():
    """Build timestamp->text index from agent_sessions.json."""
    index = []
    sessions = load_json(AGENT_SESSIONS_FILE)
    if not isinstance(sessions, list):
        return index
    for s in sessions:
        ts = _parse_ts(s.get("timestamp", ""))
        if not ts:
            continue
        parts = []
        for l in s.get("learnings", []):
            if isinstance(l, str):
                parts.append(l)
        for f in s.get("failures", []):
            if isinstance(f, str):
                parts.append(f"Failure: {f}")
        if parts:
            index.append((ts, " | ".join(parts)))
    return index


def _fuzzy_match(item_ts, log_index, window_s=FUZZY_WINDOW_S):
    """Find log entry within fuzzy window of item timestamp."""
    if not item_ts:
        return None
    for log_ts, user_text, agent_text in log_index:
        diff = abs((item_ts - log_ts).total_seconds())
        if diff <= window_s:
            desc_parts = []
            if user_text:
                desc_parts.append(f"User: {user_text[:100]}")
            if agent_text:
                desc_parts.append(f"Agent: {agent_text[:100]}")
            return " | ".join(desc_parts) if desc_parts else None
    return None


def _fuzzy_match_sessions(item_ts, session_index, window_s=FUZZY_WINDOW_S):
    """Find agent_sessions entry within fuzzy window."""
    if not item_ts:
        return None
    for sess_ts, text in session_index:
        diff = abs((item_ts - sess_ts).total_seconds())
        if diff <= window_s:
            return text[:200]
    return None


def run_backfill(dry_run=True):
    """Main backfill logic."""
    print(f"=== Backfill Session Descriptions [{'DRY-RUN' if dry_run else 'EXECUTE'}] ===\n")

    if not SESSION_FILE.exists():
        print("Error: current_session.json not found")
        return

    items = load_json(SESSION_FILE)
    if not isinstance(items, list):
        print("Error: current_session.json is not a list")
        return

    total = len(items)
    stubs = [i for i, item in enumerate(items)
             if item.get("description") in STUB_DESCRIPTIONS
             or not item.get("type")
             or not item.get("source")]
    print(f"Total items: {total}")
    print(f"Items needing fix: {len(stubs)}")

    if not stubs:
        print("Nothing to backfill!")
        return

    print("\nBuilding log index...")
    log_index = _build_log_index()
    print(f"  Log entries: {len(log_index)}")

    print("Building session index...")
    session_index = _build_session_index()
    print(f"  Session entries: {len(session_index)}")

    stats = {"log_matched": 0, "session_matched": 0, "fallback": 0, "type_fixed": 0, "source_fixed": 0}

    for idx in stubs:
        item = items[idx]
        item_ts = _parse_ts(item.get("timestamp", ""))

        # Fix missing type
        if not item.get("type"):
            item["type"] = "OBSERVATION"
            stats["type_fixed"] += 1

        # Fix missing source
        if not item.get("source"):
            item["source"] = "pulse_bridge"
            stats["source_fixed"] += 1

        # Fix stub description
        if item.get("description") in STUB_DESCRIPTIONS:
            # Try log match first
            desc = _fuzzy_match(item_ts, log_index)
            if desc:
                item["description"] = desc
                stats["log_matched"] += 1
                continue

            # Try session match
            desc = _fuzzy_match_sessions(item_ts, session_index)
            if desc:
                item["description"] = desc
                stats["session_matched"] += 1
                continue

            # Last resort
            item["description"] = "[historical - no source text]"
            stats["fallback"] += 1

    print(f"\n=== Results ===")
    print(f"  Log-matched:     {stats['log_matched']}")
    print(f"  Session-matched: {stats['session_matched']}")
    print(f"  Fallback label:  {stats['fallback']}")
    print(f"  Type fixed:      {stats['type_fixed']}")
    print(f"  Source fixed:    {stats['source_fixed']}")

    if not dry_run:
        save_json(SESSION_FILE, items)
        print(f"\nSaved to {SESSION_FILE}")
    else:
        print(f"\nDry run — no changes written. Use --execute to apply.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Backfill session descriptions")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--dry-run", action="store_true", help="Preview only")
    group.add_argument("--execute", action="store_true", help="Apply changes")
    args = parser.parse_args()
    run_backfill(dry_run=args.dry_run)
