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

"""
CDR-Aware BloomOS Loop

This is a minimal, safe adaptation of CDR_Steven_Package's bloomos/loop.py
to your current helix_cdc stack.

It performs:
    - task classification via BloomOSAgent
    - execution via BloomOSClient
    - SymCapsule wrapping
    - Se measurement
    - pheromone recording

Usage:

    PYTHONPATH=. python3 -m helix_cdc.bloomos.loop_cdr \
        --goal "solve Tower of Hanoi with 3 disks" \
        --pheromone-path /tmp/bloomos_pheromones.jsonl

This does not replace your runtime; it's an optional lab loop.
"""

import argparse
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Set

from .bloomos_client import BloomOSClient
from .bloomos_agent import BloomOSAgent, AgentStep
from .sym import SymCapsule, lower_task_to_sym
from .pheromone import PheromoneBoard, make_entry_from_step
from .gate_controller import GateController
from .program_controller import ProgramController  # Phase 7
from helix_cdc.metrics import measure_text, SeMetrics


@dataclass
class LoopResult:
    goal: str
    task_type: str
    gate_policy: 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]
    capsule: Optional[SymCapsule]
    raw: Dict[str, Any]


class CDRLoop:
    def __init__(
        self,
        base_url: str = "http://localhost:7070",
        timeout: float = 60.0,
        pheromone_path: Optional[str] = None,
        log_path: Optional[str] = None,
        adaptive_gate: bool = True,
        adaptive_program: bool = False,  # Phase 7
        max_regrow_depth: Optional[int] = None,  # Phase 8 safety
        allowed_gates: Optional[List[str]] = None,  # Phase 8 safety
        allowed_programs: Optional[List[str]] = None,  # Phase 8 safety (names or sequences)
    ) -> None:
        self.client = BloomOSClient(base_url=base_url, timeout=timeout, log_path=log_path)
        self.agent = BloomOSAgent(client=self.client)
        self.pheromones = PheromoneBoard(pheromone_path) if pheromone_path else None
        self.log_path = Path(log_path) if log_path else None
        self.adaptive_gate = adaptive_gate
        # Phase 6: GateController for adaptive gate selection
        self.gate_controller = GateController(pheromone_path) if pheromone_path and adaptive_gate else None
        # Phase 7: ProgramController for adaptive helix_program selection
        self.adaptive_program = adaptive_program
        self.program_controller = (
            ProgramController(pheromone_path) if pheromone_path and adaptive_program else None
        )
        # Phase 8: Safety rails
        self.max_regrow_depth = max_regrow_depth
        self.allowed_gates: Optional[Set[str]] = set(allowed_gates) if allowed_gates else None
        self.allowed_programs: Optional[Set[str]] = set(allowed_programs) if allowed_programs else None

    def _extract_output_text(self, raw: Dict[str, Any]) -> Optional[str]:
        out = raw.get("output")
        if out is None:
            return None
        if isinstance(out, str):
            return out
        if isinstance(out, dict):
            for key in ("text", "content", "response", "output", "result", "reasoning"):
                if key in out and isinstance(out[key], str):
                    return out[key]
            return json.dumps(out)
        return str(out)

    def run_once(self, goal: str, helix_program: Optional[str] = None) -> LoopResult:
        # 1) classify goal
        step: AgentStep = self.agent.classify_goal(goal)

        # Phase 6: Use GateController for adaptive gate selection
        gate_policy = step.gate_policy
        if self.gate_controller is not None:
            adaptive_gate = self.gate_controller.choose_gate_policy(step.task_type)
            if adaptive_gate:
                gate_policy = adaptive_gate

        # Phase 8: Enforce allowed_gates, if provided
        if self.allowed_gates is not None and gate_policy not in self.allowed_gates:
            # Fallback to 'balanced' if allowed, else first allowed gate
            if "balanced" in self.allowed_gates:
                gate_policy = "balanced"
            elif self.allowed_gates:
                gate_policy = sorted(self.allowed_gates)[0]

        # Phase 7: Use ProgramController for adaptive helix_program selection
        chosen_program: Optional[str] = helix_program
        if chosen_program is None and self.program_controller is not None:
            rec = self.program_controller.recommend_program(step.task_type)
            if rec and rec.sequence:
                chosen_program = rec.sequence

        # Phase 8: Enforce allowed_programs, if provided
        if self.allowed_programs is not None and chosen_program is not None:
            if chosen_program not in self.allowed_programs:
                # Fallback: use ABCDEFGH if allowed, else first allowed program
                if "ABCDEFGH" in self.allowed_programs:
                    chosen_program = "ABCDEFGH"
                elif "BALANCED64" in self.allowed_programs:
                    chosen_program = "ABCDEFGH"  # BALANCED64 sequence
                elif self.allowed_programs:
                    chosen_program = sorted(self.allowed_programs)[0]
                else:
                    chosen_program = None

        extra: Dict[str, Any] = dict(step.extra)
        if chosen_program is not None:
            extra["helix_program"] = chosen_program
        # Phase 8: Enforce max_regrow_depth if set
        if self.max_regrow_depth is not None:
            extra["max_regrow_depth"] = self.max_regrow_depth

        t0 = time.perf_counter()
        result = self.client.run_task(
            task_type=step.task_type,
            description=goal,
            anchors=step.anchors,
            gate_policy=gate_policy,
            extra=extra,
        )
        elapsed_ms = (time.perf_counter() - t0) * 1000.0

        raw = result.raw
        program_used = raw.get("helix_program_used") or chosen_program or "ABCDEFGH"

        # 2) Se measurement
        text_out = self._extract_output_text(raw)
        se_H = se_C = se_D = se_Se = None
        if text_out:
            try:
                m: SeMetrics = measure_text(text_out, auto_detect_d=True)
                se_H, se_C, se_D, se_Se = m.H, m.C, m.D, m.Se
            except Exception:
                pass

        # 3) Capsule
        capsule = None
        try:
            capsule = lower_task_to_sym(
                task_type=step.task_type,
                description=goal,
                program=program_used,
                result=raw,
            )
        except Exception:
            pass

        # 4) Pheromones
        if self.pheromones is not None:
            try:
                gate_used = raw.get("gate_used") or step.gate_policy
                entry = make_entry_from_step(
                    task_type=step.task_type,
                    gate=gate_used,
                    helix_program=program_used,
                    status=result.status,
                    logits_mean=result.logits_mean,
                    logits_std=result.logits_std,
                    se_H=se_H,
                    se_C=se_C,
                    se_D=se_D,
                    se_Se=se_Se,
                    meta={
                        "goal": goal,
                        "elapsed_ms": elapsed_ms,
                    },
                )
                self.pheromones.record(entry)
            except Exception:
                pass

        loop_result = LoopResult(
            goal=goal,
            task_type=step.task_type,
            gate_policy=step.gate_policy,
            helix_program=helix_program,
            status=result.status,
            logits_mean=result.logits_mean,
            logits_std=result.logits_std,
            se_H=se_H,
            se_C=se_C,
            se_D=se_D,
            se_Se=se_Se,
            capsule=capsule,
            raw=raw,
        )

        # Optional logging
        if self.log_path is not None:
            try:
                entry = {
                    "timestamp": time.time(),
                    "goal": goal,
                    "task_type": step.task_type,
                    "gate_policy": step.gate_policy,
                    "helix_program": helix_program,
                    "status": result.status,
                    "logits_mean": result.logits_mean,
                    "logits_std": result.logits_std,
                    "se_metrics": {
                        "H": se_H,
                        "C": se_C,
                        "D": se_D,
                        "Se": se_Se,
                    },
                    "capsule": capsule.to_dict() if capsule else None,
                }
                with self.log_path.open("a", encoding="utf-8") as f:
                    f.write(json.dumps(entry) + "\n")
            except Exception:
                pass

        return loop_result


def _main() -> None:
    parser = argparse.ArgumentParser(
        prog="loop_cdr",
        description="CDR-aware BloomOS loop (single-step).",
    )
    parser.add_argument(
        "--goal",
        required=True,
        help="Natural language goal/command to run.",
    )
    parser.add_argument(
        "--base-url",
        default="http://localhost:7070",
        help="BloomOS runtime base URL.",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=60.0,
        help="HTTP timeout for BloomOS calls.",
    )
    parser.add_argument(
        "--helix-program",
        help="Optional helix_program override (named or raw 8-codon sequence).",
    )
    parser.add_argument(
        "--log-path",
        help="Optional JSONL log path for loop entries.",
    )
    parser.add_argument(
        "--pheromone-path",
        help="Optional JSONL pheromone board path.",
    )
    parser.add_argument(
        "--adaptive-program",
        action="store_true",
        help="Enable adaptive helix_program selection based on pheromones (Phase 7).",
    )
    # Phase 8: Safety rails
    parser.add_argument(
        "--max-regrow-depth",
        type=int,
        help="Maximum allowed regrow depth for CDR seed expansion (Phase 8 safety).",
    )
    parser.add_argument(
        "--allowed-gate",
        action="append",
        dest="allowed_gates",
        help="Allowed gate policy (can be specified multiple times).",
    )
    parser.add_argument(
        "--allowed-program",
        action="append",
        dest="allowed_programs",
        help="Allowed program name or sequence (can be specified multiple times).",
    )
    args = parser.parse_args()

    loop = CDRLoop(
        base_url=args.base_url,
        timeout=args.timeout,
        pheromone_path=args.pheromone_path,
        log_path=args.log_path,
        adaptive_gate=True,
        adaptive_program=args.adaptive_program,
        max_regrow_depth=args.max_regrow_depth,
        allowed_gates=args.allowed_gates,
        allowed_programs=args.allowed_programs,
    )
    res = loop.run_once(goal=args.goal, helix_program=args.helix_program)

    print("=== CDR Loop Result ===")
    print(f"goal:         {res.goal}")
    print(f"task_type:    {res.task_type}")
    print(f"gate_policy:  {res.gate_policy}")
    print(f"status:       {res.status}")
    print(f"logits_mean:  {res.logits_mean}")
    print(f"Se:           {res.se_Se}")
    if res.capsule:
        print(f"capsule soul: {res.capsule.soul.get('soul_id', '')[:16]}...")


if __name__ == "__main__":
    _main()
