# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
D17: Predictive Swarm Choreographer — Concurrency Optimization.

Reads lock contention data from brain_io instrumentation, identifies
hot-file clusters, and generates choreography rules that stagger
agent execution to prevent I/O contention.

Source: DMN session_20260227T033543, Gemini 3.1 Pro
Builds on: skill_dependency_synthesizer, conflict_resolver, brain_io, procedural_tracker
"""

import json
import logging
import time
from datetime import datetime, timezone
from pathlib import Path

_log = logging.getLogger("pulse.d17_choreographer")

try:
    from pulse.core.brain_utils import STATE_DIR
except Exception:
    STATE_DIR = Path("state")

CONTENTION_FILE = STATE_DIR / "health" / "contention_tracker.json"
CHOREOGRAPHY_FILE = STATE_DIR / "health" / "choreography_rules.json"

# Thresholds
_MIN_WAITS_FOR_HOT = 3         # file needs 3+ contention events to be "hot"
_MIN_AVG_WAIT_S = 1.0          # average wait > 1s = meaningful contention
_STAGGER_BASE_MS = 200         # base stagger delay between agents
_MAX_RULES = 30                # cap choreography rules


def _load_contention() -> dict:
    """Load contention tracker data from brain_io instrumentation."""
    if not CONTENTION_FILE.exists():
        return {}
    try:
        with open(CONTENTION_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return {}


def _identify_hot_files(contention: dict) -> list[dict]:
    """Identify files with significant lock contention."""
    hot = []
    for fname, stats in contention.items():
        total = stats.get("total_waits", 0)
        total_s = stats.get("total_seconds", 0.0)
        if total < _MIN_WAITS_FOR_HOT:
            continue
        avg_wait = total_s / total if total > 0 else 0
        if avg_wait < _MIN_AVG_WAIT_S:
            continue
        hot.append({
            "file": fname,
            "total_waits": total,
            "avg_wait": round(avg_wait, 3),
            "max_wait": stats.get("max_wait", 0.0),
            "recent": stats.get("recent", [])[-5:],
        })
    hot.sort(key=lambda x: -x["total_waits"])
    return hot


def _detect_co_contention(contention: dict) -> list[dict]:
    """Detect files that experience contention in the same time windows.

    If task_board.json and agent_registry.json both spike within 10s,
    they form a co-contention cluster — agents touching one likely
    touch the other, creating cascading lock waits.
    """
    # Build time-bucketed events per file (10s buckets)
    buckets: dict[int, list[str]] = {}
    for fname, stats in contention.items():
        for event in stats.get("recent", []):
            ts = event.get("t", "")
            try:
                dt = datetime.fromisoformat(ts)
                bucket = int(dt.timestamp()) // 10
                buckets.setdefault(bucket, []).append(fname)
            except (ValueError, TypeError):
                continue

    # Find files that co-occur in the same time buckets
    co_pairs: dict[str, int] = {}
    for files in buckets.values():
        unique = sorted(set(files))
        for i, a in enumerate(unique):
            for b in unique[i + 1:]:
                key = f"{a}|{b}"
                co_pairs[key] = co_pairs.get(key, 0) + 1

    clusters = []
    for pair, count in sorted(co_pairs.items(), key=lambda x: -x[1]):
        if count < 2:
            break
        a, b = pair.split("|")
        clusters.append({"files": [a, b], "co_occurrences": count})

    return clusters[:10]


def _generate_rules(hot_files: list, clusters: list) -> list[dict]:
    """Generate choreography rules from contention analysis."""
    rules = []
    now = datetime.now(timezone.utc).isoformat()

    # Rule type 1: Stagger access to hot files
    for hf in hot_files[:10]:
        stagger_ms = min(
            int(hf["avg_wait"] * 1000 * 0.5),  # half the avg wait
            5000  # cap at 5s
        )
        stagger_ms = max(stagger_ms, _STAGGER_BASE_MS)
        rules.append({
            "type": "stagger",
            "file": hf["file"],
            "stagger_ms": stagger_ms,
            "reason": f"avg wait {hf['avg_wait']:.1f}s across {hf['total_waits']} events",
            "created": now,
        })

    # Rule type 2: Sequence co-contention clusters
    for cluster in clusters[:5]:
        rules.append({
            "type": "sequence",
            "files": cluster["files"],
            "order": cluster["files"],  # lock in alphabetical order to prevent deadlock
            "reason": f"{cluster['co_occurrences']} co-contention events",
            "created": now,
        })

    return rules[:_MAX_RULES]


def run_once(verbose: bool = False) -> dict:
    """Run one choreography analysis cycle.

    Returns summary dict with hot_files, clusters, and rules generated.
    """
    contention = _load_contention()
    if not contention:
        if verbose:
            _log.info("No contention data yet — brain_io hasn't recorded any waits > 500ms")
        return {"hot_files": 0, "clusters": 0, "rules": 0, "status": "no_data"}

    hot_files = _identify_hot_files(contention)
    clusters = _detect_co_contention(contention)
    rules = _generate_rules(hot_files, clusters)

    if verbose:
        _log.info("Hot files: %d, Co-contention clusters: %d, Rules: %d",
                  len(hot_files), len(clusters), len(rules))

    # Save choreography rules
    if rules:
        output = {
            "generated": datetime.now(timezone.utc).isoformat(),
            "hot_files": hot_files[:10],
            "clusters": clusters[:5],
            "rules": rules,
            "contention_summary": {
                fname: {"waits": s.get("total_waits", 0),
                        "max": s.get("max_wait", 0)}
                for fname, s in contention.items()
            },
        }
        CHOREOGRAPHY_FILE.parent.mkdir(parents=True, exist_ok=True)
        tmp = CHOREOGRAPHY_FILE.with_suffix(".tmp")
        try:
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(output, f, indent=2, ensure_ascii=False)
            import os
            os.replace(str(tmp), str(CHOREOGRAPHY_FILE))
        except OSError:
            if tmp.exists():
                tmp.unlink()

    return {
        "hot_files": len(hot_files),
        "clusters": len(clusters),
        "rules": len(rules),
        "status": "ok",
    }


def get_stagger_ms(filename: str) -> int:
    """Query: get recommended stagger delay for a file.

    Called by task claiming code to add pre-lock delay,
    spreading agent access across time to reduce contention.
    Returns 0 if no rule exists.
    """
    if not CHOREOGRAPHY_FILE.exists():
        return 0
    try:
        with open(CHOREOGRAPHY_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        for rule in data.get("rules", []):
            if rule.get("type") == "stagger" and rule.get("file") == filename:
                return rule.get("stagger_ms", 0)
    except (json.JSONDecodeError, OSError):
        pass
    return 0


def get_lock_order(files: list[str]) -> list[str]:
    """Query: get recommended lock acquisition order for a set of files.

    Prevents deadlocks by ensuring all agents acquire locks in the
    same order when touching co-contention file clusters.
    Returns sorted list (alphabetical as safe default).
    """
    if not CHOREOGRAPHY_FILE.exists():
        return sorted(files)
    try:
        with open(CHOREOGRAPHY_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        for rule in data.get("rules", []):
            if rule.get("type") == "sequence":
                rule_files = set(rule.get("files", []))
                if rule_files.intersection(files):
                    # Use the rule's order for known files, append unknowns
                    ordered = [f for f in rule["order"] if f in files]
                    ordered += sorted(f for f in files if f not in ordered)
                    return ordered
    except (json.JSONDecodeError, OSError):
        pass
    return sorted(files)


if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    verbose = "--verbose" in sys.argv or "-v" in sys.argv
    result = run_once(verbose=True)
    print(json.dumps(result, indent=2))
