#!/usr/bin/env python3
"""
helix_cdc/fibpi_runtime/runtime.py

WO-FIBPI-DNA-01:
Define a FibPiRuntime that acts as the "cell/runtime" for DNA seeds.
For v1 this can be mostly wiring + stubs; later WOs will swap in real growth logic.
"""

from __future__ import annotations

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

import numpy as np


@dataclass
class FibPiRuntimeConfig:
    """Configuration for the FibPi runtime."""
    dim: int = 3
    num_bands: int = 8
    seed_dir: Path = Path("seeds")
    base_model_gguf: Optional[Path] = None  # for v1 "cheat" growth


@dataclass
class FibPiAnchor:
    """A location on the FibPi lattice."""
    coords: List[float]  # phi-lattice coordinates
    band_id: int         # entropy/frequency band
    dims: int            # lattice dimensionality


@dataclass
class FibPiDNALatent:
    """Latent code associated with a block's DNA."""
    code_helix64: str
    dim: int


@dataclass
class FibPiDNASeed:
    """Full DNA seed record for one block."""
    model_family: str
    block_id: str
    anchor: FibPiAnchor
    latent: FibPiDNALatent
    decoder_hints: Dict[str, Any]


class FibPiRuntime:
    """
    FibPi runtime: interprets DNA seeds in the context of the phi-lattice environment
    and produces tensor blocks.
    """

    def __init__(self, config: FibPiRuntimeConfig):
        self.config = config
        # TODO: load FibPi lattice, NF1 invariants, etc.
        self._seed_cache: Dict[str, FibPiDNASeed] = {}
        # For v1, optionally hold a reference to a GGUF reader:
        self._base_reader = None

    # --- Seed management -----------------------------------------------------

    def load_seed_for_block(self, block_id: str) -> Optional[FibPiDNASeed]:
        """
        Load DNA seed for the given block_id from disk or cache.

        v1: look for seeds/{block_id}.fibpi_dna.json (or similar schema).
        """
        if block_id in self._seed_cache:
            return self._seed_cache[block_id]

        # v1: filename convention seeds/{block_id}.fibpi_dna.json
        # Replace dots with underscores for filename safety
        safe_block_id = block_id.replace(".", "_")
        candidate = self.config.seed_dir / f"{safe_block_id}.fibpi_dna.json"
        if not candidate.exists():
            return None

        with candidate.open() as f:
            raw = json.load(f)

        anchor = FibPiAnchor(
            coords=raw["anchor"]["fibpi_coord"],
            band_id=int(raw["anchor"]["band_id"]),
            dims=int(raw["anchor"]["dims"]),
        )
        latent = FibPiDNALatent(
            code_helix64=raw["latent"]["code"],
            dim=int(raw["latent"]["dim"]),
        )
        seed = FibPiDNASeed(
            model_family=raw["model_family"],
            block_id=raw["block_id"],
            anchor=anchor,
            latent=latent,
            decoder_hints=raw.get("decoder_hints", {}),
        )
        self._seed_cache[block_id] = seed
        return seed

    def place_seed(
        self,
        block_id: str,
        anchor: FibPiAnchor,
        band_id: int,
        latent: FibPiDNALatent,
    ) -> None:
        """
        Register a seed at a given anchor/band. In v1 this may just
        populate an in-memory structure or write a JSON schema file.
        """
        seed = FibPiDNASeed(
            model_family="unknown",
            block_id=block_id,
            anchor=anchor,
            latent=latent,
            decoder_hints={},
        )
        self._seed_cache[block_id] = seed

    # --- Growth --------------------------------------------------------------

    def grow_block(self, block_id: str) -> np.ndarray:
        """
        Grow a tensor block from the FibPi lattice + DNA seed.

        v1: honesty mode:
        - Load seed (for wiring)
        - Return the original tensor from base GGUF
        v2+: substitute learned / procedural growth.
        """
        seed = self.load_seed_for_block(block_id)
        if seed is None:
            raise KeyError(f"No DNA seed found for block_id={block_id}")

        # v1: just "cheat" by loading from GGUF so KAT can compare.
        print(f"[FIBPI-RUNTIME] Growing block {block_id} using seed at anchor={seed.anchor}")
        block = self._load_block_from_base_model(block_id)
        return block

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

    def _ensure_base_reader(self):
        """Lazy-load GGUF reader for the base model, if configured."""
        if self._base_reader is not None:
            return
        if self.config.base_model_gguf is None:
            raise RuntimeError("base_model_gguf not configured on FibPiRuntime")
        from gguf import GGUFReader  # type: ignore
        self._base_reader = GGUFReader(self.config.base_model_gguf)

    def _load_block_from_base_model(self, block_id: str) -> np.ndarray:
        """Load tensor block from the base GGUF as float32 or quant bytes."""
        self._ensure_base_reader()
        assert self._base_reader is not None
        # Find tensor by name
        for tensor in self._base_reader.tensors:
            if tensor.name == block_id:
                return np.array(tensor.data).astype(np.float32)
        raise KeyError(f"Tensor {block_id} not found in base model")
