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

"""
Phase 7: Adaptive Program Controller

This module selects helix_programs based on pheromone history.
It uses:
    - PheromoneBoard entries (with Se, seed_entropy, lattice_D, etc.)
    - fitness() from fitness.py
    - program definitions from helix_programs.json via program_loader
    - GATE_SEMANTICS from programs.py for fallback

Primary entrypoint:

    ProgramController.recommend_program(task_type: str) -> ProgramChoice

This is complementary to GateController:
    - GateController chooses gate_policy
    - ProgramController chooses helix_program
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

import json
import math
import time

from .pheromone import PheromoneBoard, PheromoneEntry
from .fitness import fitness
from .program_loader import resolve_program, list_programs, get_program_for_task
from .programs import GATE_SEMANTICS


@dataclass
class ProgramChoice:
    """Recommendation from the ProgramController."""
    name: Optional[str]
    sequence: str
    score: float
    samples: int


class ProgramController:
    """
    Adaptive program selector based on pheromone history.

    Strategy:
        - Look at last N pheromone entries for a task_type
        - Group by helix_program
        - Compute average fitness per program
        - Require at least min_samples
        - Return best program if available, else fall back to task-specific defaults
    """

    def __init__(
        self,
        pheromone_path: Optional[str] = None,
        window: int = 500,
        min_samples: int = 3,
    ) -> None:
        """
        Initialize the program controller.

        Args:
            pheromone_path: Path to pheromone JSONL file
            window: Number of most recent entries to consider
            min_samples: Minimum samples required for a program to be considered
        """
        self.pheromone_path = Path(pheromone_path) if pheromone_path else None
        self.window = window
        self.min_samples = min_samples

    # --------------------------
    # Internal helpers
    # --------------------------

    def _iter_recent_entries(self) -> Iterable[PheromoneEntry]:
        """Load and iterate over recent pheromone entries."""
        if not self.pheromone_path or not self.pheromone_path.exists():
            return []
        board = PheromoneBoard(str(self.pheromone_path))
        # PheromoneBoard.iter_entries yields in file order; we want last N
        entries = list(board.iter_entries())
        if self.window > 0 and len(entries) > self.window:
            entries = entries[-self.window:]
        return entries

    def _group_scores(
        self, entries: Iterable[PheromoneEntry], task_type: str
    ) -> Dict[str, Tuple[float, int]]:
        """
        Group pheromone entries by helix_program (sequence string), compute:
            - total fitness
            - sample count

        Returns:
            dict of program_sequence -> (total_fitness, count)
        """
        scores: Dict[str, Tuple[float, int]] = {}
        for e in entries:
            if e.task_type != task_type:
                continue
            if not e.helix_program:
                continue
            try:
                f = fitness(e)
            except Exception:
                continue
            total, count = scores.get(e.helix_program, (0.0, 0))
            scores[e.helix_program] = (total + f, count + 1)
        return scores

    def _map_sequence_to_name(self, seq: str) -> Optional[str]:
        """
        Try to map a program sequence back to a named program, using:
            - helix_programs.json via program_loader.list_programs()
            - GATE_SEMANTICS from programs.py
        """
        # Named programs from helix_programs.json
        try:
            named = list_programs()
            for name, s in named.items():
                if s == seq:
                    return name
        except Exception:
            pass

        # Gate semantics (legacy gate-based programs)
        try:
            for gate_name, info in GATE_SEMANTICS.items():
                if info.get("program") == seq:
                    return f"gate:{gate_name}"
        except Exception:
            pass

        return None

    # --------------------------
    # Public API
    # --------------------------

    def recommend_program(self, task_type: str) -> Optional[ProgramChoice]:
        """
        Recommend a helix_program for the given task_type, based on pheromone history.

        Strategy:
            - Look at last N pheromone entries for this task_type
            - Group by helix_program
            - Compute average fitness per program
            - Require at least min_samples
            - Return best program if available, else fall back to task defaults

        Args:
            task_type: The task type (e.g., "math", "code", "json", "reason")

        Returns:
            ProgramChoice with name, sequence, score, and sample count
        """
        entries = list(self._iter_recent_entries())
        if not entries:
            return self._fallback_program(task_type, reason="no_pheromones")

        scores = self._group_scores(entries, task_type)
        if not scores:
            return self._fallback_program(task_type, reason="no_matching_task")

        best_seq = None
        best_score = -math.inf
        best_count = 0

        for seq, (total, count) in scores.items():
            if count < self.min_samples:
                continue
            avg = total / max(count, 1)
            if avg > best_score:
                best_score = avg
                best_seq = seq
                best_count = count

        if best_seq is None:
            return self._fallback_program(task_type, reason="insufficient_samples")

        name = self._map_sequence_to_name(best_seq)
        return ProgramChoice(
            name=name,
            sequence=best_seq,
            score=best_score,
            samples=best_count,
        )

    def _fallback_program(
        self, task_type: str, reason: str = "fallback"
    ) -> Optional[ProgramChoice]:
        """
        Fallback to task-specific program when pheromones are not informative.

        Uses:
            1. get_program_for_task() from program_loader (helix_programs.json)
            2. Ultimate fallback: BALANCED64 (ABCDEFGH)
        """
        try:
            # Try task-specific program from helix_programs.json
            seq, name = get_program_for_task(task_type)
            return ProgramChoice(
                name=name,
                sequence=seq,
                score=0.0,
                samples=0,
            )
        except Exception:
            pass

        # Ultimate fallback: BALANCED64
        return ProgramChoice(
            name="BALANCED64",
            sequence="ABCDEFGH",
            score=0.0,
            samples=0,
        )

    def get_all_recommendations(self) -> Dict[str, Optional[ProgramChoice]]:
        """
        Get recommendations for all task_types found in pheromone history.

        Returns:
            Dict mapping task_type -> ProgramChoice
        """
        entries = list(self._iter_recent_entries())
        task_types = set(e.task_type for e in entries if e.task_type)

        recommendations = {}
        for tt in task_types:
            recommendations[tt] = self.recommend_program(tt)

        return recommendations

    def inspect(self) -> Dict[str, Any]:
        """
        Inspect the pheromone history and return analysis.

        Returns:
            Dict with program usage statistics and recommendations
        """
        if self.pheromone_path is None or not self.pheromone_path.exists():
            return {"error": "No pheromone file", "path": str(self.pheromone_path)}

        entries = list(self._iter_recent_entries())

        # Count programs per task
        counts: Dict[str, Dict[str, int]] = {}
        for e in entries:
            if not e.helix_program:
                continue
            tt = e.task_type or "unknown"
            prog = e.helix_program
            counts.setdefault(tt, {})
            counts[tt][prog] = counts[tt].get(prog, 0) + 1

        # Get recommendations
        recommendations = {}
        for tt in counts.keys():
            rec = self.recommend_program(tt)
            if rec:
                recommendations[tt] = {
                    "name": rec.name,
                    "sequence": rec.sequence,
                    "score": round(rec.score, 4),
                    "samples": rec.samples,
                }

        return {
            "path": str(self.pheromone_path),
            "total_entries": len(entries),
            "window": self.window,
            "min_samples": self.min_samples,
            "program_usage": counts,
            "recommendations": recommendations,
        }


# ------------------------------
# CLI for inspection
# ------------------------------

def _cli_inspect(path: str, task_type: Optional[str]) -> None:
    """CLI inspection handler."""
    controller = ProgramController(pheromone_path=path)

    if task_type:
        rec = controller.recommend_program(task_type)
        data = {
            "task_type": task_type,
            "recommendation": {
                "name": rec.name if rec else None,
                "sequence": rec.sequence if rec else None,
                "score": round(rec.score, 4) if rec else None,
                "samples": rec.samples if rec else None,
            }
            if rec
            else None,
        }
    else:
        data = controller.inspect()

    print(json.dumps(data, indent=2))


def _main() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="program_controller",
        description="Phase 7: Inspect or recommend helix_programs based on pheromones.",
    )
    parser.add_argument(
        "--inspect",
        metavar="PATH",
        required=True,
        help="Path to pheromone JSONL file.",
    )
    parser.add_argument(
        "--task-type",
        help="If provided, show recommendation for this task_type.",
    )
    parser.add_argument(
        "--window",
        type=int,
        default=500,
        help="Number of recent entries to consider (default: 500).",
    )
    parser.add_argument(
        "--min-samples",
        type=int,
        default=3,
        help="Minimum samples required for selection (default: 3).",
    )
    args = parser.parse_args()

    controller = ProgramController(
        pheromone_path=args.inspect,
        window=args.window,
        min_samples=args.min_samples,
    )

    if args.task_type:
        rec = controller.recommend_program(args.task_type)
        data = {
            "task_type": args.task_type,
            "recommendation": {
                "name": rec.name if rec else None,
                "sequence": rec.sequence if rec else None,
                "score": round(rec.score, 4) if rec else None,
                "samples": rec.samples if rec else None,
            }
            if rec
            else None,
        }
        print(json.dumps(data, indent=2))
    else:
        data = controller.inspect()
        print(json.dumps(data, indent=2))


if __name__ == "__main__":
    _main()
