#!/usr/bin/env python3
"""
Session Doctor — Automatic Error Recovery for AI Agent Sessions
Core library module.

Usage as CLI:
    python -m session_doctor monitor              # Start monitoring daemon
    python -m session_doctor status               # Show health dashboard
    python -m session_doctor inject <session_id> <error_type>  # Test injection
    python -m session_doctor report               # Show recovery history
    python -m session_doctor demo                 # Run self-contained demo

Usage as library:
    from session_doctor import SessionDoctor, Config
    doctor = SessionDoctor()
    doctor.register_session("my-session", label="coding-agent")
    doctor.ingest_event("my-session", "Error: 429 rate limit exceeded")
"""

import os
import re
import sys
import time
import json
import sqlite3
import argparse
import threading
from enum import Enum
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


# ─── Path resolution ──────────────────────────────────────────────────────────
# Priority order:
#   1. SESSION_DOCTOR_HOME env var (explicit override)
#   2. OpenClaw legacy path (~/.openclaw/workspace/projects/session-doctor/) if it exists
#   3. Default: ~/.session-doctor/

def _resolve_base_dir() -> Path:
    env_home = os.environ.get("SESSION_DOCTOR_HOME")
    if env_home:
        return Path(env_home).expanduser()
    legacy = Path.home() / ".openclaw" / "workspace" / "projects" / "session-doctor"
    if legacy.exists():
        return legacy
    return Path.home() / ".session-doctor"


_BASE_DIR = _resolve_base_dir()
DEFAULT_DB_PATH = _BASE_DIR / "state.db"
DEFAULT_LOG_PATH = _BASE_DIR / "session_doctor.log"

ERROR_PATTERNS = {
    "http_400":    re.compile(r"(400|Bad Request|invalid_request)", re.IGNORECASE),
    "token_limit": re.compile(r"(context.?length|token.?limit|max.?tokens|too.?long)", re.IGNORECASE),
    "json_parse":  re.compile(r"(JSON|parse.?error|invalid.?json|SyntaxError)", re.IGNORECASE),
    "timeout":     re.compile(r"(timeout|timed.?out|ETIMEDOUT|deadline)", re.IGNORECASE),
    "tool_repeat": re.compile(r"REPEATED_TOOL_CALL", re.IGNORECASE),
    "rate_limit":  re.compile(r"(429|rate.?limit|too.?many.?requests)", re.IGNORECASE),
    "auth_error":  re.compile(r"(401|403|Unauthorized|Forbidden)", re.IGNORECASE),
    "budget_over": re.compile(r"(budget.?exceeded|over.?budget|cost.?limit)", re.IGNORECASE),
}


class Severity(Enum):
    LOW      = 1
    MEDIUM   = 2
    HIGH     = 3
    CRITICAL = 4


class Strategy(Enum):
    MONITOR           = "monitor"
    BACKOFF           = "backoff"
    RETRY_PROMPT_FIX  = "retry_prompt_fix"
    COMPACT_CONTEXT   = "compact_context"
    BREAK_LOOP        = "break_loop"
    RESET             = "reset"
    HALT              = "halt"


class SessionStatus(Enum):
    OK      = "ok"
    WARN    = "warn"
    ERROR   = "error"
    HALTED  = "halted"


# Maps error type → (Severity, Strategy)
ERROR_POLICY = {
    "http_400":    (Severity.HIGH,     Strategy.RESET),
    "token_limit": (Severity.HIGH,     Strategy.COMPACT_CONTEXT),
    "json_parse":  (Severity.MEDIUM,   Strategy.RETRY_PROMPT_FIX),
    "timeout":     (Severity.MEDIUM,   Strategy.BACKOFF),
    "tool_repeat": (Severity.HIGH,     Strategy.BREAK_LOOP),
    "rate_limit":  (Severity.LOW,      Strategy.BACKOFF),
    "auth_error":  (Severity.CRITICAL, Strategy.HALT),
    "budget_over": (Severity.CRITICAL, Strategy.HALT),
}

# ─── Config ───────────────────────────────────────────────────────────────────

@dataclass
class Config:
    check_interval_sec: int = 30
    error_repeat_count: int = 3
    error_window_sec: int = 300
    idle_timeout_sec: int = 90
    budget_warn_pct: float = 0.80
    budget_cap: float = 2.00
    max_tool_depth: int = 10
    max_auto_retries: int = 2
    compact_keep_last_n: int = 5
    backoff_base_sec: float = 2.0
    backoff_max_sec: float = 32.0
    notifications_enabled: bool = True
    notifications_min_severity: Severity = Severity.MEDIUM


# ─── Data Models ──────────────────────────────────────────────────────────────

@dataclass
class ErrorEvent:
    session_id: str
    error_type: str
    error_detail: str
    detected_at: datetime = field(default_factory=datetime.now)
    severity: Severity = Severity.LOW
    strategy: Strategy = Strategy.MONITOR


@dataclass
class SessionState:
    session_id: str
    label: str = ""
    status: SessionStatus = SessionStatus.OK
    started_at: datetime = field(default_factory=datetime.now)
    last_seen: datetime = field(default_factory=datetime.now)
    budget_used: float = 0.0
    error_count: int = 0
    recent_errors: list = field(default_factory=list)
    recovery_count: int = 0


# ─── State Store ──────────────────────────────────────────────────────────────

class StateStore:
    def __init__(self, db_path: Path = DEFAULT_DB_PATH):
        db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(str(db_path), check_same_thread=False)
        self._lock = threading.Lock()
        self._init_schema()

    def _init_schema(self):
        with self.conn:
            self.conn.executescript("""
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    label TEXT DEFAULT '',
                    status TEXT DEFAULT 'ok',
                    started_at TEXT,
                    last_seen TEXT,
                    budget_used REAL DEFAULT 0.0,
                    error_count INTEGER DEFAULT 0
                );

                CREATE TABLE IF NOT EXISTS errors (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT,
                    error_type TEXT,
                    error_detail TEXT,
                    detected_at TEXT,
                    severity TEXT,
                    strategy TEXT,
                    outcome TEXT DEFAULT 'pending'
                );

                CREATE TABLE IF NOT EXISTS recoveries (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT,
                    strategy TEXT,
                    triggered_at TEXT,
                    completed_at TEXT,
                    success INTEGER DEFAULT 0,
                    notes TEXT DEFAULT ''
                );
            """)

    def upsert_session(self, s: SessionState):
        with self._lock:
            with self.conn:
                self.conn.execute("""
                    INSERT INTO sessions (session_id, label, status, started_at, last_seen, budget_used, error_count)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(session_id) DO UPDATE SET
                        label=excluded.label,
                        status=excluded.status,
                        last_seen=excluded.last_seen,
                        budget_used=excluded.budget_used,
                        error_count=excluded.error_count
                """, (
                    s.session_id, s.label, s.status.value,
                    s.started_at.isoformat(), s.last_seen.isoformat(),
                    s.budget_used, s.error_count
                ))

    def log_error(self, e: ErrorEvent) -> int:
        with self._lock:
            with self.conn:
                cur = self.conn.execute("""
                    INSERT INTO errors (session_id, error_type, error_detail, detected_at, severity, strategy)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (
                    e.session_id, e.error_type, e.error_detail,
                    e.detected_at.isoformat(), e.severity.name, e.strategy.name
                ))
                return cur.lastrowid

    def update_error_outcome(self, error_id: int, outcome: str):
        with self._lock:
            with self.conn:
                self.conn.execute(
                    "UPDATE errors SET outcome=? WHERE id=?",
                    (outcome, error_id)
                )

    def log_recovery(self, session_id: str, strategy: str, success: bool, notes: str = "") -> int:
        with self._lock:
            with self.conn:
                cur = self.conn.execute("""
                    INSERT INTO recoveries (session_id, strategy, triggered_at, completed_at, success, notes)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (
                    session_id, strategy,
                    datetime.now().isoformat(), datetime.now().isoformat(),
                    1 if success else 0, notes
                ))
                return cur.lastrowid

    def get_all_sessions(self):
        with self._lock:
            cur = self.conn.execute("SELECT * FROM sessions ORDER BY last_seen DESC")
            return cur.fetchall()

    def get_recent_errors(self, session_id: str, window_sec: int = 300):
        cutoff = (datetime.now() - timedelta(seconds=window_sec)).isoformat()
        with self._lock:
            cur = self.conn.execute("""
                SELECT error_type, COUNT(*) as cnt
                FROM errors
                WHERE session_id=? AND detected_at > ?
                GROUP BY error_type
            """, (session_id, cutoff))
            return {row[0]: row[1] for row in cur.fetchall()}

    def get_recovery_history(self, limit: int = 20):
        with self._lock:
            cur = self.conn.execute("""
                SELECT r.*, s.label FROM recoveries r
                LEFT JOIN sessions s ON r.session_id = s.session_id
                ORDER BY r.triggered_at DESC LIMIT ?
            """, (limit,))
            return cur.fetchall()


# ─── Detector ─────────────────────────────────────────────────────────────────

class Detector:
    """Classifies raw text/log lines into typed error events."""

    def classify(self, text: str) -> Optional[str]:
        """Return the first matching error type, or None."""
        for error_type, pattern in ERROR_PATTERNS.items():
            if pattern.search(text):
                return error_type
        return None

    def is_stuck(self, session: SessionState, config: Config, store: StateStore) -> Optional[str]:
        """
        Returns the dominant error type if session is stuck, else None.
        A session is 'stuck' if the same error type appears >= threshold times
        within the error window.
        """
        recent = store.get_recent_errors(session.session_id, config.error_window_sec)
        for error_type, count in recent.items():
            if count >= config.error_repeat_count:
                return error_type

        # Check idle timeout
        idle_secs = (datetime.now() - session.last_seen).total_seconds()
        if idle_secs > config.idle_timeout_sec and session.status == SessionStatus.OK:
            return "timeout"

        return None


# ─── Recoverer ────────────────────────────────────────────────────────────────

class Recoverer:
    """Executes recovery strategies. All methods are logged and idempotent."""

    def __init__(self, store: StateStore, config: Config, notifier=None):
        self.store = store
        self.config = config
        self.notifier = notifier
        self._backoff_counts: dict[str, int] = {}

    def recover(self, session: SessionState, error_type: str, error_id: int) -> bool:
        severity, strategy = ERROR_POLICY.get(error_type, (Severity.LOW, Strategy.MONITOR))

        _log(f"[{session.session_id}] Attempting recovery: {strategy.name} (severity={severity.name})")

        if session.recovery_count >= self.config.max_auto_retries and severity != Severity.CRITICAL:
            _log(f"[{session.session_id}] Max retries ({self.config.max_auto_retries}) reached, escalating to HALT")
            strategy = Strategy.HALT

        success = False
        notes = ""

        if strategy == Strategy.MONITOR:
            success = True
            notes = "No action needed (LOW severity)"
        elif strategy == Strategy.BACKOFF:
            success, notes = self._do_backoff(session)
        elif strategy == Strategy.COMPACT_CONTEXT:
            success, notes = self._do_compact_context(session)
        elif strategy == Strategy.RETRY_PROMPT_FIX:
            success, notes = self._do_retry_prompt_fix(session)
        elif strategy == Strategy.BREAK_LOOP:
            success, notes = self._do_break_loop(session)
        elif strategy == Strategy.RESET:
            success, notes = self._do_reset(session)
        elif strategy == Strategy.HALT:
            success, notes = self._do_halt(session)

        self.store.log_recovery(session.session_id, strategy.name, success, notes)
        self.store.update_error_outcome(error_id, "recovered" if success else "failed")

        if self.notifier and self.config.notifications_enabled:
            if severity.value >= self.config.notifications_min_severity.value:
                self.notifier.send(session, error_type, strategy, success, notes)

        return success

    def _do_backoff(self, session: SessionState) -> tuple[bool, str]:
        count = self._backoff_counts.get(session.session_id, 0)
        wait = min(
            self.config.backoff_base_sec * (2 ** count),
            self.config.backoff_max_sec
        )
        self._backoff_counts[session.session_id] = count + 1
        _log(f"[{session.session_id}] BACKOFF: waiting {wait:.0f}s (attempt {count+1})")
        time.sleep(wait)
        return True, f"Backed off {wait:.0f}s (attempt {count+1})"

    def _do_compact_context(self, session: SessionState) -> tuple[bool, str]:
        _log(f"[{session.session_id}] COMPACT_CONTEXT: trimming to last {self.config.compact_keep_last_n} tool results")
        notes = f"Simulated compact: kept last {self.config.compact_keep_last_n} entries"
        return True, notes

    def _do_retry_prompt_fix(self, session: SessionState) -> tuple[bool, str]:
        _log(f"[{session.session_id}] RETRY_PROMPT_FIX: injecting JSON format reminder")
        notes = "Injected format correction prompt (simulated)"
        return True, notes

    def _do_break_loop(self, session: SessionState) -> tuple[bool, str]:
        _log(f"[{session.session_id}] BREAK_LOOP: injecting stop-loop event")
        notes = "Loop broken via stop_loop system event (simulated)"
        return True, notes

    def _do_reset(self, session: SessionState) -> tuple[bool, str]:
        _log(f"[{session.session_id}] RESET: saving state snapshot, restarting session")
        snapshot_path = DEFAULT_DB_PATH.parent / f"snapshot_{session.session_id}_{int(time.time())}.json"
        snapshot = {
            "session_id": session.session_id,
            "label": session.label,
            "status": session.status.value,
            "budget_used": session.budget_used,
            "error_count": session.error_count,
            "snapshot_at": datetime.now().isoformat(),
        }
        snapshot_path.write_text(json.dumps(snapshot, indent=2))
        _log(f"[{session.session_id}] Snapshot saved to {snapshot_path}")
        notes = f"State snapshot saved to {snapshot_path.name}. Restart simulated."
        return True, notes

    def _do_halt(self, session: SessionState) -> tuple[bool, str]:
        _log(f"[{session.session_id}] HALT: stopping session, requires manual review")
        session.status = SessionStatus.HALTED
        self.store.upsert_session(session)
        notes = "Session halted. Manual intervention required."
        return True, notes


# ─── Notifier ─────────────────────────────────────────────────────────────────

class Notifier:
    """
    Sends structured alerts. Logs to console + writes to file.
    Real impl would call OpenClaw message tool to Telegram.
    """

    def __init__(self, output_path: Optional[Path] = None):
        self.output_path = output_path or (DEFAULT_DB_PATH.parent / "notifications.log")

    def send(self, session: SessionState, error_type: str, strategy: Strategy,
             success: bool, notes: str):
        status_emoji = "✅" if success else "❌"
        msg = (
            f"\n🏥 Session Doctor Alert\n"
            f"{'─' * 40}\n"
            f"Session : {session.session_id} ({session.label or 'unlabeled'})\n"
            f"Issue   : {error_type} (count: {session.error_count})\n"
            f"Action  : {strategy.name}\n"
            f"Status  : {status_emoji} {'Recovered' if success else 'Failed — manual review needed'}\n"
            f"Notes   : {notes}\n"
            f"Budget  : ${session.budget_used:.2f} / $2.00\n"
            f"Time    : {datetime.now().strftime('%Y-%m-%d %H:%M:%S PST')}\n"
            f"{'─' * 40}"
        )
        _log(f"NOTIFICATION: {msg}")
        with open(self.output_path, "a") as f:
            f.write(msg + "\n")


# ─── Session Doctor ───────────────────────────────────────────────────────────

class SessionDoctor:
    """
    Main entry point. Register sessions, feed events, let Session Doctor
    detect and recover from errors automatically.

    Example::

        doctor = SessionDoctor()
        doctor.register_session("my-run", label="coding-agent")
        doctor.ingest_event("my-run", "Error: 429 rate limit exceeded")
    """

    def __init__(self, config: Config = None):
        self.config = config or Config()
        self.store = StateStore()
        self.detector = Detector()
        self.notifier = Notifier()
        self.recoverer = Recoverer(self.store, self.config, self.notifier)
        self._sessions: dict[str, SessionState] = {}
        self._running = False

    def register_session(self, session_id: str, label: str = "") -> SessionState:
        """Register a new session for monitoring."""
        s = SessionState(session_id=session_id, label=label)
        self._sessions[session_id] = s
        self.store.upsert_session(s)
        _log(f"Registered session: {session_id} ({label})")
        return s

    def ingest_event(self, session_id: str, text: str, budget_used: float = 0.0):
        """
        Feed a log line or event text into Session Doctor.
        Call this whenever your session emits output.

        :param session_id: Unique identifier for the session.
        :param text: Raw text to classify (log line, API response, etc.).
        :param budget_used: Current cumulative cost in USD (optional).
        """
        session = self._sessions.get(session_id)
        if not session:
            session = self.register_session(session_id)

        session.last_seen = datetime.now()
        session.budget_used = max(session.budget_used, budget_used)

        error_type = self.detector.classify(text)
        if error_type:
            session.error_count += 1
            severity, strategy = ERROR_POLICY.get(error_type, (Severity.LOW, Strategy.MONITOR))
            event = ErrorEvent(
                session_id=session_id,
                error_type=error_type,
                error_detail=text[:200],
                severity=severity,
                strategy=strategy,
            )
            error_id = self.store.log_error(event)

            if severity == Severity.CRITICAL:
                session.status = SessionStatus.ERROR
            elif severity == Severity.HIGH and session.status == SessionStatus.OK:
                session.status = SessionStatus.WARN

            self.store.upsert_session(session)

            stuck_type = self.detector.is_stuck(session, self.config, self.store)
            if stuck_type:
                _log(f"[{session_id}] Stuck detected on '{stuck_type}' — triggering recovery")
                session.recovery_count += 1
                self.recoverer.recover(session, stuck_type, error_id)

        if session.budget_used >= self.config.budget_cap:
            _log(f"[{session_id}] Budget exceeded! ${session.budget_used:.2f}")
            session.status = SessionStatus.ERROR
            event = ErrorEvent(
                session_id=session_id,
                error_type="budget_over",
                error_detail=f"Budget ${session.budget_used:.2f} >= cap ${self.config.budget_cap:.2f}",
                severity=Severity.CRITICAL,
                strategy=Strategy.HALT,
            )
            error_id = self.store.log_error(event)
            self.recoverer.recover(session, "budget_over", error_id)

    def check_all_sessions(self):
        """Periodic health check across all registered sessions."""
        for session_id, session in list(self._sessions.items()):
            stuck_type = self.detector.is_stuck(session, self.config, self.store)
            if stuck_type and session.status != SessionStatus.HALTED:
                _log(f"[{session_id}] Periodic check: stuck on '{stuck_type}'")
                error_event = ErrorEvent(
                    session_id=session_id,
                    error_type=stuck_type,
                    error_detail="Detected during periodic check",
                )
                error_id = self.store.log_error(error_event)
                session.recovery_count += 1
                self.recoverer.recover(session, stuck_type, error_id)

    def run(self):
        """Start the main monitoring loop (blocking)."""
        self._running = True
        _log(f"Session Doctor started (interval={self.config.check_interval_sec}s)")
        try:
            while self._running:
                self.check_all_sessions()
                time.sleep(self.config.check_interval_sec)
        except KeyboardInterrupt:
            _log("Session Doctor stopped by user")
        finally:
            self._running = False

    def stop(self):
        """Stop the monitoring loop."""
        self._running = False


# ─── Dashboard ────────────────────────────────────────────────────────────────

STATUS_EMOJI = {
    "ok":     "🟢",
    "warn":   "🟡",
    "error":  "🔴",
    "halted": "⚫",
}


def print_dashboard(store: StateStore):
    sessions = store.get_all_sessions()
    print()
    print("╔══════════════════════════════════════════════════════════════╗")
    print("║             Session Doctor — Health Dashboard                ║")
    print("╠══════════════════╦════════════╦══════════╦═════════╦════════╣")
    print("║ Session          ║ Status     ║ Errors   ║ Budget  ║ Label  ║")
    print("╠══════════════════╬════════════╬══════════╬═════════╬════════╣")

    if not sessions:
        print("║  No sessions registered yet.                                 ║")
    else:
        for row in sessions:
            sid, label, status, started, last_seen, budget, errors = row[:7]
            emoji = STATUS_EMOJI.get(status, "❓")
            sid_display = sid[:16].ljust(16)
            status_display = f"{emoji} {status}".ljust(10)
            errors_display = str(errors).ljust(8)
            budget_display = f"${budget:.2f}".ljust(7)
            label_display = (label or "-")[:6].ljust(6)
            print(f"║ {sid_display} ║ {status_display} ║ {errors_display} ║ {budget_display} ║ {label_display} ║")

    print("╚══════════════════╩════════════╩══════════╩═════════╩════════╝")
    print(f"\nLast check: {datetime.now().strftime('%Y-%m-%d %H:%M:%S PST')}")
    print()


def print_report(store: StateStore):
    rows = store.get_recovery_history(limit=20)
    print("\n=== Recovery History (last 20) ===\n")
    if not rows:
        print("  No recoveries recorded yet.")
        return
    for row in rows:
        rid, sid, strategy, triggered, completed, success, notes, label = row
        outcome = "✅" if success else "❌"
        print(f"  {outcome} [{triggered[:16]}] {sid} | {strategy} | {notes[:60]}")
    print()


# ─── Logging ──────────────────────────────────────────────────────────────────

def _log(msg: str):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line)
    DEFAULT_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(DEFAULT_LOG_PATH, "a") as f:
        f.write(line + "\n")


# ─── CLI ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Session Doctor — Automatic Error Recovery for AI Agent Sessions"
    )
    sub = parser.add_subparsers(dest="command")

    sub.add_parser("monitor", help="Start monitoring daemon")
    sub.add_parser("status", help="Show health dashboard")
    sub.add_parser("report", help="Show recovery history")

    inject_p = sub.add_parser("inject", help="Inject a test error event")
    inject_p.add_argument("session_id", help="Session ID to target")
    inject_p.add_argument("error_type", choices=list(ERROR_PATTERNS.keys()), help="Error type to inject")

    sub.add_parser("demo", help="Run a self-contained demo")

    args = parser.parse_args()

    if args.command == "status":
        store = StateStore()
        print_dashboard(store)

    elif args.command == "report":
        store = StateStore()
        print_report(store)

    elif args.command == "inject":
        doctor = SessionDoctor()
        for i in range(3):
            text = f"Simulated error: {args.error_type} (occurrence {i+1})"
            _log(f"Injecting: {text}")
            doctor.ingest_event(args.session_id, text)
            time.sleep(0.1)

    elif args.command == "monitor":
        doctor = SessionDoctor()
        doctor.register_session("heartbeat", label="heartbeat")
        doctor.run()

    elif args.command == "demo":
        print("\n🏥 Session Doctor — Demo Run\n")
        doctor = SessionDoctor()

        doctor.register_session("session-abc123", label="gh-issues")
        doctor.register_session("session-def456", label="coding-agent")

        print("\n--- Simulating rate limit loop on session-abc123 ---")
        for i in range(4):
            doctor.ingest_event("session-abc123", f"Error: 429 rate limit exceeded (retry {i})")
            time.sleep(0.1)

        print("\n--- Simulating token limit loop on session-def456 ---")
        for i in range(4):
            doctor.ingest_event("session-def456", f"Error: context length exceeded max_tokens (attempt {i})")
            time.sleep(0.1)

        print_dashboard(doctor.store)
        print_report(doctor.store)

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
