#!/usr/bin/env python3
"""
HelixBackplane: FibPi3D → Position-4 Gate → Helix64 VM

Thin wrapper around existing pieces with FlowTorch-style braiding:
- gates.py (resolve_anchors, choose_gate, apply_gate_to_program)
- bridge.py (_execute_via_helix64_http, _execute_via_bloomos_kernel)
"""

from __future__ import annotations

import os
import time
from typing import Dict, Any, List, Optional, Tuple

from .gates import (
    resolve_anchors,
    choose_gate,
    apply_gate_to_program,
    GATE_POLICIES,
)
from .bridge import (
    _execute_via_helix64_http,
    _execute_via_bloomos_kernel,
)
from .programs import (
    BASE_PROGRAM,
    GATE_POSITION,
    program_to_ops,
    describe_program,
    GATE_SEMANTICS,
)
from .program_loader import (
    resolve_program,
    is_valid_codon_sequence,
)


class HelixBackplane:
    """
    Unified execution surface for Helix64 tasks.

    Responsibilities:
      - Resolve anchors from task_type/overrides
      - Select gate policy (math_bias/code_bias/json_bias/balanced)
      - Mutate program at position 4
      - Execute via chosen backend (direct/http or kernel)
      - Return structured result with metadata
    """

    def __init__(
        self,
        base_program: str = None,  # Uses BASE_PROGRAM from programs.py
        default_backend: str = "helix64",  # helix64 | kernel | hybrid
        default_budget_ms: float = 100.0,
        default_fallback: str = "helix64",  # helix64 | stub | zero
    ) -> None:
        self.base_program = base_program or BASE_PROGRAM
        self.default_backend = default_backend
        self.default_budget_ms = default_budget_ms
        self.default_fallback = default_fallback

    def _get_backend_config(self) -> Tuple[str, float, str]:
        backend = os.getenv("BLOOMOS_BACKEND", self.default_backend).lower()
        budget_ms = float(os.getenv("BLOOMOS_BUDGET_MS", str(self.default_budget_ms)))
        fallback = os.getenv("BLOOMOS_FALLBACK", self.default_fallback).lower()
        return backend, budget_ms, fallback

    def run_task(
        self,
        task_type: str,
        description: str,
        anchors: Optional[List[str]] = None,
        gate_policy: Optional[str] = None,
        base_program: Optional[str] = None,
        helix_program: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Main entrypoint.

        Returns:
          {
            "task_type": ...,
            "anchors_resolved": [...],
            "gate": {
              "policy": ...,
              "codon": ...,
              "program_before": ...,
              "program_after": ...,
              "expected_effects": {...},
            },
            "backend_results": {...},
            "logits_mean": float,
            "logits_std": float,
            "meta": {...},
        }
        """
        backend, budget_ms, fallback = self._get_backend_config()

        # Resolve helix_program if provided (takes precedence over base_program)
        helix_program_resolved = None
        helix_program_name = None
        if helix_program:
            helix_program_resolved, helix_program_name = resolve_program(helix_program)

        program_before = helix_program_resolved or base_program or self.base_program
        meta: Dict[str, Any] = {
            "backend": backend,
            "budget_ms": budget_ms,
            "fallback_policy": fallback,
        }

        # 1) Anchors + gate
        anchors_resolved = resolve_anchors(task_type, anchors)
        gate_policy_name, gate_cfg = choose_gate(task_type, gate_policy)

        # If full helix_program provided, use it directly (no gate mutation)
        # Otherwise apply gate mutation to position 4
        if helix_program_resolved:
            program_after = helix_program_resolved  # Use full program as-is
            meta["helix_program_name"] = helix_program_name
            meta["helix_program_mode"] = "full_override"
        else:
            program_after = apply_gate_to_program(program_before, gate_cfg)
            meta["helix_program_mode"] = "gate_mutated"

        meta["gate_policy"] = gate_policy_name
        meta["program_before"] = program_before
        meta["program_after"] = program_after

        # 2) Execute with time budget + fallbacks
        t0 = time.perf_counter()
        result: Optional[Dict[str, Any]] = None
        used_fallback = False
        backend_error: Optional[str] = None

        try:
            if backend == "kernel":
                # Future path – BloomOS kernel mode
                result = _execute_via_bloomos_kernel(
                    program_after, anchors_resolved, task_type, description
                )

            elif backend == "hybrid":
                # Try kernel, fall back to direct HTTP
                try:
                    result = _execute_via_bloomos_kernel(
                        program_after, anchors_resolved, task_type, description
                    )
                except Exception as e:
                    backend_error = f"kernel error: {str(e)[:200]}"
                    result = _execute_via_helix64_http(program_after, anchors_resolved)
                    used_fallback = True

            else:
                # Default: direct HTTP to Helix64 VM
                result = _execute_via_helix64_http(program_after, anchors_resolved)

        except Exception as e:
            backend_error = f"backend exception: {str(e)[:200]}"
            result = None

        elapsed_ms = (time.perf_counter() - t0) * 1000.0
        meta["elapsed_ms"] = round(elapsed_ms, 3)

        # 3) Budget check + fallback
        if (result is None) or (elapsed_ms > budget_ms):
            used_fallback = True
            if fallback == "helix64":
                result = _execute_via_helix64_http(program_after, anchors_resolved)
            elif fallback == "stub":
                # Deterministic stub
                result = {
                    a: {"status": "ok", "logits_mean": 0.5, "logits_std": 0.1}
                    for a in anchors_resolved
                }
            else:  # "zero"
                result = {
                    a: {"status": "ok", "logits_mean": 0.0, "logits_std": 0.0}
                    for a in anchors_resolved
                }

        meta["used_fallback"] = used_fallback
        if backend_error:
            meta["backend_error"] = backend_error

        # 4) Aggregate logits
        logits_mean_list: List[float] = []
        logits_std_list: List[float] = []

        for a, r in (result or {}).items():
            if r.get("status") == "ok":
                if "logits_mean" in r:
                    logits_mean_list.append(r["logits_mean"])
                if "logits_std" in r:
                    logits_std_list.append(r["logits_std"])

        agg_mean = sum(logits_mean_list) / len(logits_mean_list) if logits_mean_list else 0.0
        agg_std = sum(logits_std_list) / len(logits_std_list) if logits_std_list else 0.0

        # Get semantic descriptions
        ops_before = program_to_ops(program_before)
        ops_after = program_to_ops(program_after)
        gate_semantics = GATE_SEMANTICS.get(gate_policy_name, {})

        return {
            "task_type": task_type,
            "description": description,
            "anchors_resolved": anchors_resolved,
            "gate": {
                "policy": gate_policy_name,
                "codon": gate_cfg["codon"],
                "program_before": program_before,
                "program_after": program_after,
                "ops_before": ops_before,
                "ops_after": ops_after,
                "gate_position": GATE_POSITION,
                "gate_op_before": ops_before[GATE_POSITION] if len(ops_before) > GATE_POSITION else None,
                "gate_op_after": ops_after[GATE_POSITION] if len(ops_after) > GATE_POSITION else None,
                "expected_effects": gate_cfg["expected_effects"],
                "mutation_description": gate_semantics.get("effect", ""),
            },
            "backend_results": result,
            "logits_mean": agg_mean,
            "logits_std": agg_std,
            "meta": meta,
        }


# Convenience singleton
DEFAULT_BACKPLANE = HelixBackplane()
