#!/usr/bin/env python3
from __future__ import annotations

"""
Pheromone Board for BloomOS / Helix64.

This is a lightweight feedback store:
    - Each "pheromone" entry describes how good a configuration was:
        (task_type, gate, helix_program, Se, logits_mean, success)
    - Stored as JSONL.
    - We can later query which gates/programs tend to perform best for a task_type.

This is intentionally simple and file-based.
"""

from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Union
import json
import time


@dataclass
class PheromoneEntry:
    timestamp: float
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]
    status: str
    logits_mean: Optional[float]
    logits_std: Optional[float]
    se_H: Optional[float]
    se_C: Optional[float]
    se_D: Optional[float]
    se_Se: Optional[float]
    meta: Dict[str, Any]
    # Phase 5 enriched fields for adaptive learning
    seed_len: Optional[int] = None
    seed_entropy: Optional[float] = None
    regrow_depth: Optional[int] = None
    lattice_D: Optional[int] = None
    # Kernel-mode IT token tracking (NEW)
    it_tokens_spent: Optional[float] = None
    syscall_used: Optional[str] = None


class PheromoneBoard:
    """
    File-backed pheromone store.

    Typical usage:
        board = PheromoneBoard("/tmp/bloomos_pheromones.jsonl")
        board.record(...)
        best = board.best_for_task("reason")
    """

    def __init__(self, path: Union[str, Path]) -> None:
        self.path = Path(path)

    def record(self, entry: PheromoneEntry) -> None:
        rec = asdict(entry)
        try:
            with self.path.open("a", encoding="utf-8") as f:
                f.write(json.dumps(rec) + "\n")
        except OSError:
            pass  # non-fatal

    def iter_entries(self) -> Iterable[PheromoneEntry]:
        if not self.path.exists():
            return
        with self.path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    data = json.loads(line)
                    yield PheromoneEntry(**data)
                except Exception:
                    continue

    def best_for_task(
        self,
        task_type: str,
        metric: str = "se_Se",
        min_success_rate: float = 0.0,
    ) -> Optional[Dict[str, Any]]:
        """
        Find the best (gate, helix_program) pair for a given task_type
        according to a metric (default: Se).

        Returns:
            {
                "gate": ...,
                "helix_program": ...,
                "metric_avg": ...,
                "success_rate": ...,
                "num_samples": ...,
            }
        or None if no data.
        """
        # Aggregate: (gate, prog) -> list[entries]
        buckets: Dict[tuple, List[PheromoneEntry]] = {}
        for e in self.iter_entries():
            if e.task_type != task_type:
                continue
            key = (e.gate, e.helix_program)
            buckets.setdefault(key, []).append(e)

        best_key = None
        best_value: Optional[float] = None
        best_info: Optional[Dict[str, Any]] = None

        for (gate, prog), entries in buckets.items():
            vals = []
            ok = 0
            for e in entries:
                if metric == "se_Se":
                    v = e.se_Se
                elif metric == "logits_mean":
                    v = e.logits_mean
                else:
                    v = getattr(e, metric, None)
                if v is not None:
                    vals.append(float(v))
                if e.status == "ok":
                    ok += 1

            if not vals:
                continue

            avg = sum(vals) / len(vals)
            success_rate = ok / len(entries) if entries else 0.0
            if success_rate < min_success_rate:
                continue

            if best_value is None or avg > best_value:
                best_value = avg
                best_key = (gate, prog)
                best_info = {
                    "gate": gate,
                    "helix_program": prog,
                    "metric_avg": avg,
                    "success_rate": success_rate,
                    "num_samples": len(entries),
                }

        return best_info

    def all_for_task(
        self,
        task_type: str,
        metric: str = "se_Se",
    ) -> List[Dict[str, Any]]:
        """
        Get all (gate, helix_program) configurations ranked by metric for task_type.

        Returns list of dicts sorted by metric_avg descending.
        """
        buckets: Dict[tuple, List[PheromoneEntry]] = {}
        for e in self.iter_entries():
            if e.task_type != task_type:
                continue
            key = (e.gate, e.helix_program)
            buckets.setdefault(key, []).append(e)

        results = []
        for (gate, prog), entries in buckets.items():
            vals = []
            ok = 0
            for e in entries:
                if metric == "se_Se":
                    v = e.se_Se
                elif metric == "logits_mean":
                    v = e.logits_mean
                else:
                    v = getattr(e, metric, None)
                if v is not None:
                    vals.append(float(v))
                if e.status == "ok":
                    ok += 1

            if not vals:
                continue

            avg = sum(vals) / len(vals)
            success_rate = ok / len(entries) if entries else 0.0
            results.append({
                "gate": gate,
                "helix_program": prog,
                "metric_avg": avg,
                "success_rate": success_rate,
                "num_samples": len(entries),
            })

        results.sort(key=lambda x: x["metric_avg"], reverse=True)
        return results


def make_entry_from_step(
    task_type: str,
    gate: Optional[str],
    helix_program: Optional[str],
    status: str,
    logits_mean: Optional[float],
    logits_std: Optional[float],
    se_H: Optional[float],
    se_C: Optional[float],
    se_D: Optional[float],
    se_Se: Optional[float],
    meta: Optional[Dict[str, Any]] = None,
    # Phase 5 enriched fields
    seed_len: Optional[int] = None,
    seed_entropy: Optional[float] = None,
    regrow_depth: Optional[int] = None,
    lattice_D: Optional[int] = None,
    # Kernel-mode IT token tracking (NEW)
    it_tokens_spent: Optional[float] = None,
    syscall_used: Optional[str] = None,
) -> PheromoneEntry:
    return PheromoneEntry(
        timestamp=time.time(),
        task_type=task_type,
        gate=gate,
        helix_program=helix_program,
        status=status,
        logits_mean=logits_mean,
        logits_std=logits_std,
        se_H=se_H,
        se_C=se_C,
        se_D=se_D,
        se_Se=se_Se,
        meta=meta or {},
        seed_len=seed_len,
        seed_entropy=seed_entropy,
        regrow_depth=regrow_depth,
        lattice_D=lattice_D,
        it_tokens_spent=it_tokens_spent,
        syscall_used=syscall_used,
    )
