#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Procedural Tracker: Track event sequences and success rates.

Records which procedures agents follow (e.g., brain_query -> Edit -> success)
and tracks their success rates over time. High-success procedures become
recommended workflows.

Usage:
    python procedural_tracker.py --record brain_query_before_edit --success
    python procedural_tracker.py --record brain_query_before_edit --failure
    python procedural_tracker.py --stats brain_query_before_edit
    python procedural_tracker.py --recommend architecture

Dependencies: Python stdlib + brain_utils (internal)
"""
import argparse
import io
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

if os.name == 'nt' and __name__ == '__main__':
    if hasattr(sys.stdout, 'reconfigure'):
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
    if hasattr(sys.stderr, 'reconfigure'):
        sys.stderr.reconfigure(encoding='utf-8', errors='replace')


from pulse.core.brain_utils import (
    HIPPOCAMPUS_DIR, load_json_dict, save_json, atomic_update_json,
    now_iso as _now_iso,
)

PROCEDURAL_LOG = HIPPOCAMPUS_DIR / "Evidence" / "Metrics" / "procedural_log.json"

# Known procedures and their descriptions
_KNOWN_PROCEDURES = {
    "brain_query_before_edit": {
        "name": "Brain Query Before Edit",
        "description": "Query brain before file modifications",
        "event_sequence": ["brain_query", "Edit"],
        "domain": "architecture",
    },
    "brain_query_before_implement": {
        "name": "Brain Query Before Implementation",
        "description": "Query brain before starting any implementation task",
        "event_sequence": ["brain_query", "implement"],
        "domain": "architecture",
    },
    "test_after_change": {
        "name": "Test After Change",
        "description": "Run tests after modifying code",
        "event_sequence": ["Edit", "test"],
        "domain": "testing",
    },
    "save_after_task": {
        "name": "Save After Task Completion",
        "description": "Run pulse_save after completing a task",
        "event_sequence": ["complete_task", "pulse_save"],
        "domain": "operations",
    },
    "anti_pattern_check": {
        "name": "Anti-Pattern Check",
        "description": "Check anti-patterns before implementing",
        "event_sequence": ["check_anti_patterns", "implement"],
        "domain": "security",
    },
}


def _load_procedures():
    """Load procedure log, initializing if needed."""
    data = load_json_dict(PROCEDURAL_LOG)
    if "procedures" not in data:
        data["procedures"] = []
    return data


def _find_procedure(data, procedure_id):
    """Find a procedure by ID in the log."""
    for p in data["procedures"]:
        if p.get("procedure_id") == procedure_id:
            return p
    return None


def record_execution(procedure_id, success=True, duration_seconds=0.0):
    """Record a procedure execution (success or failure).

    Args:
        procedure_id: Procedure identifier (e.g., "brain_query_before_edit")
        success: Whether the procedure succeeded
        duration_seconds: How long the procedure took

    Returns:
        Updated procedure stats dict
    """
    def _updater(data):
        if not isinstance(data, dict):
            data = {"procedures": []}
        if "procedures" not in data:
            data["procedures"] = []

        # Find or create procedure entry
        proc = None
        for p in data["procedures"]:
            if p.get("procedure_id") == procedure_id:
                proc = p
                break

        if proc is None:
            # Create new entry from known procedures or generic
            known = _KNOWN_PROCEDURES.get(procedure_id, {})
            proc = {
                "procedure_id": procedure_id,
                "name": known.get("name", procedure_id.replace("_", " ").title()),
                "description": known.get("description", ""),
                "event_sequence": known.get("event_sequence", []),
                "executions": 0,
                "successes": 0,
                "failures": 0,
                "success_rate": 0.0,
                "avg_duration_seconds": 0.0,
                "domain": known.get("domain", "general"),
                "first_recorded": _now_iso(),
                "last_executed": None,
                "confidence": 0.0,
            }
            data["procedures"].append(proc)

        # Update stats
        proc["executions"] += 1
        if success:
            proc["successes"] += 1
        else:
            proc["failures"] += 1

        total = proc["executions"]
        proc["success_rate"] = round(proc["successes"] / total, 3) if total > 0 else 0.0
        proc["last_executed"] = _now_iso()

        # Update running average duration
        if duration_seconds > 0:
            old_avg = proc.get("avg_duration_seconds", 0.0)
            proc["avg_duration_seconds"] = round(
                (old_avg * (total - 1) + duration_seconds) / total, 2
            )

        # Confidence = success_rate weighted by sample size (bayesian-ish)
        # Small samples have low confidence, large samples converge
        proc["confidence"] = round(
            proc["success_rate"] * (1 - 1 / (total + 1)), 3
        )

        data["last_updated"] = _now_iso()
        return data

    result = atomic_update_json(PROCEDURAL_LOG, _updater, default={"procedures": []})
    proc = _find_procedure(result, procedure_id)
    return proc or {}


def get_procedure_stats(procedure_id):
    """Get stats for a specific procedure.

    Returns:
        Procedure dict with success_rate, confidence, etc. Empty dict if not found.
    """
    data = _load_procedures()
    proc = _find_procedure(data, procedure_id)
    return proc or {}


def recommend_procedures(domain="general", min_confidence=0.5):
    """Recommend proven procedures for a given domain.

    Args:
        domain: Task domain (e.g., "architecture", "testing")
        min_confidence: Minimum confidence threshold

    Returns:
        List of procedure dicts sorted by confidence (highest first)
    """
    data = _load_procedures()
    candidates = []

    for proc in data.get("procedures", []):
        # Match domain or "general" procedures apply everywhere
        proc_domain = proc.get("domain", "general")
        if proc_domain != domain and proc_domain != "general" and domain != "general":
            continue

        if proc.get("confidence", 0) >= min_confidence:
            candidates.append(proc)

    # Sort by confidence descending
    candidates.sort(key=lambda p: -p.get("confidence", 0))
    return candidates


def get_all_stats():
    """Get summary statistics across all procedures."""
    data = _load_procedures()
    procedures = data.get("procedures", [])

    if not procedures:
        return {"total_procedures": 0, "total_executions": 0}

    total_exec = sum(p.get("executions", 0) for p in procedures)
    total_success = sum(p.get("successes", 0) for p in procedures)
    high_conf = [p for p in procedures if p.get("confidence", 0) >= 0.7]

    return {
        "total_procedures": len(procedures),
        "total_executions": total_exec,
        "overall_success_rate": round(total_success / total_exec, 3) if total_exec > 0 else 0,
        "high_confidence_count": len(high_conf),
        "last_updated": data.get("last_updated"),
    }


def main():
    parser = argparse.ArgumentParser(description="PULSE Procedural Tracker")
    parser.add_argument("--record", help="Record procedure execution (procedure_id)")
    parser.add_argument("--success", action="store_true", help="Mark as success (default)")
    parser.add_argument("--failure", action="store_true", help="Mark as failure")
    parser.add_argument("--duration", type=float, default=0.0, help="Duration in seconds")
    parser.add_argument("--stats", help="Show stats for a procedure")
    parser.add_argument("--recommend", help="Recommend procedures for domain")
    parser.add_argument("--all", action="store_true", help="Show all procedure stats")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    args = parser.parse_args()

    if args.record:
        success = not args.failure  # Default to success unless --failure
        result = record_execution(args.record, success=success,
                                  duration_seconds=args.duration)
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            status = "SUCCESS" if success else "FAILURE"
            print(f"[PROCEDURE] {args.record}: {status}")
            print(f"  Executions: {result.get('executions', 0)}")
            print(f"  Success rate: {result.get('success_rate', 0):.1%}")
            print(f"  Confidence: {result.get('confidence', 0):.3f}")

    elif args.stats:
        result = get_procedure_stats(args.stats)
        if not result:
            print(f"No data for procedure: {args.stats}")
        elif args.json:
            print(json.dumps(result, indent=2))
        else:
            print(f"[PROCEDURE] {result.get('name', args.stats)}")
            print(f"  Executions: {result.get('executions', 0)}")
            print(f"  Success rate: {result.get('success_rate', 0):.1%}")
            print(f"  Confidence: {result.get('confidence', 0):.3f}")
            print(f"  Domain: {result.get('domain', 'general')}")
            print(f"  Last executed: {(result.get('last_executed') or 'never')[:19]}")

    elif args.recommend:
        results = recommend_procedures(args.recommend)
        if not results:
            print(f"No proven procedures for domain: {args.recommend}")
        elif args.json:
            print(json.dumps(results, indent=2))
        else:
            print(f"[PROCEDURES] Recommended for '{args.recommend}':")
            for p in results[:5]:
                rate = p.get("success_rate", 0)
                conf = p.get("confidence", 0)
                print(f"  [{conf:.2f}] {p.get('name', '?')} "
                      f"({rate:.0%} success, {p.get('executions', 0)} runs)")

    elif args.all:
        stats = get_all_stats()
        if args.json:
            print(json.dumps(stats, indent=2))
        else:
            print(f"[PROCEDURES] Summary:")
            print(f"  Total procedures: {stats.get('total_procedures', 0)}")
            print(f"  Total executions: {stats.get('total_executions', 0)}")
            print(f"  Overall success: {stats.get('overall_success_rate', 0):.1%}")
            print(f"  High confidence: {stats.get('high_confidence_count', 0)}")

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
