#!/usr/bin/env python3
from __future__ import annotations
"""
Superglyph <-> Helix-CDC Bridge
-------------------------------

This module is the "architecture injection" point that lets the Helix-CDC
substrate become aware of the Superglyph system living elsewhere on the box.

It does *not* try to fully control the Superglyph runtime yet. Instead it:

- Describes where the Superglyph components live (genome, lab, agentic cells)
- Checks which layers are present (regen / agentic / swarm)
- Exposes a simple, typed description of the Superglyph substrate
- Provides light helpers that other modules (substrate_inspector, etc.) can use

On-box, Claude Code can then:
- Fill in TODOs for tighter integration (e.g. calling agentic cells directly)
- Hook these descriptors into controllers and dashboards.

Usage:
    PYTHONPATH=. python3 -m helix_cdc.superglyph_bridge
    PYTHONPATH=. python3 -m helix_cdc.superglyph_bridge --json
"""

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


# ---------------------------------------------------------------------------
# Constants / defaults
# ---------------------------------------------------------------------------


DEFAULT_SUPERGLYPH_ROOT = Path(os.environ.get("SUPERGLYPH_ROOT", str(Path.home() / "superglyph_lab")))
DEFAULT_GENOME_ROOT = Path(os.environ.get("SUPERGLYPH_GENOME_ROOT", str(Path.home() / "archive" / "fused_capsules")))
DEFAULT_AGENTIC_ROOT = Path(os.environ.get("AGENTIC_CELLS_ROOT", str(Path.home() / "agentic_cells")))
DEFAULT_CDR_PACKAGE_ROOT = Path(os.environ.get("CDR_PACKAGE_ROOT", str(Path.home() / "CDR_Steven_Package")))


@dataclass
class SuperglyphPaths:
    """Where things live on disk."""
    lab_root: Path
    genome_root: Path
    agentic_root: Path
    cdr_root: Path

    @classmethod
    def detect(cls) -> "SuperglyphPaths":
        return cls(
            lab_root=DEFAULT_SUPERGLYPH_ROOT,
            genome_root=DEFAULT_GENOME_ROOT,
            agentic_root=DEFAULT_AGENTIC_ROOT,
            cdr_root=DEFAULT_CDR_PACKAGE_ROOT,
        )

    def to_dict(self) -> Dict[str, str]:
        return {
            "lab_root": str(self.lab_root),
            "genome_root": str(self.genome_root),
            "agentic_root": str(self.agentic_root),
            "cdr_root": str(self.cdr_root),
        }


@dataclass
class SuperglyphCapabilities:
    """Which layers of the Superglyph stack appear to be present."""
    has_genome: bool
    has_lab_tools: bool
    has_agentic_cells: bool
    has_swarm_scripts: bool
    has_regen_core: bool

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class SuperglyphGenomeInfo:
    """Basic metadata about the main Superglyph genome (.hxz)."""
    path: Optional[str] = None
    size_bytes: Optional[int] = None
    num_lobes: Optional[int] = None
    notes: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class SuperglyphSubstrateDescriptor:
    """
    High-level descriptor of the Superglyph substrate as seen from Helix-CDC.

    This is what we "inject" into the Helix substrate as the physical symbol layer.
    """
    name: str
    version: str
    paths: SuperglyphPaths
    capabilities: SuperglyphCapabilities
    genome: SuperglyphGenomeInfo

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "version": self.version,
            "paths": self.paths.to_dict(),
            "capabilities": self.capabilities.to_dict(),
            "genome": self.genome.to_dict(),
        }


# ---------------------------------------------------------------------------
# Detection helpers
# ---------------------------------------------------------------------------


def _detect_genome(genome_root: Path) -> SuperglyphGenomeInfo:
    """
    Try to locate the primary Superglyph genome (.hxz) and derive simple stats.
    We keep this deliberately conservative: no parsing of hxz internals here.
    """
    if not genome_root.exists():
        return SuperglyphGenomeInfo(notes="genome_root not found")

    candidates: List[Path] = sorted(
        [p for p in genome_root.glob("*.hxz") if p.is_file()],
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    if not candidates:
        return SuperglyphGenomeInfo(notes="no .hxz genomes found")
    main = candidates[0]
    try:
        size = main.stat().st_size
    except OSError:
        size = None

    # We *know* from your previous receipts that Superglyph has ~28 lobes,
    # but we won't hard-code that as fact here; we just hint it.
    return SuperglyphGenomeInfo(
        path=str(main),
        size_bytes=size,
        num_lobes=None,
        notes="Main Superglyph genome candidate (lobe count not introspected here).",
    )


def _detect_capabilities(paths: SuperglyphPaths) -> SuperglyphCapabilities:
    """Check which known components of the Superglyph architecture are present."""
    lab_tools = paths.lab_root / "tools"
    regen_cli = lab_tools / "regrow_cli.sh"
    regen_kat = lab_tools / "regrow_kat.py"

    has_genome = paths.genome_root.exists()
    has_lab_tools = lab_tools.exists() and (regen_cli.exists() or regen_kat.exists())
    has_agentic_cells = paths.agentic_root.exists()
    has_swarm_scripts = (paths.agentic_root / "scripts" / "swarm_stress.sh").exists()
    helix_core_dir = paths.cdr_root / "code" / "helix_core"
    helix_regrow = helix_core_dir / "helix_regrow.py"
    has_regen_core = helix_core_dir.exists() and helix_regrow.exists()

    return SuperglyphCapabilities(
        has_genome=has_genome,
        has_lab_tools=has_lab_tools,
        has_agentic_cells=has_agentic_cells,
        has_swarm_scripts=has_swarm_scripts,
        has_regen_core=has_regen_core,
    )


# ---------------------------------------------------------------------------
# Tensor resolution helpers
# ---------------------------------------------------------------------------


def resolve_core_tensors(superglyph_manifest: Dict[str, Any], plan_path: Optional[Path] = None) -> Dict[str, Any]:
    """
    Resolve core tensor descriptors (embeddings, lm_head) from plan file.

    Args:
        superglyph_manifest: Loaded superglyph.json dict
        plan_path: Path to plan file (default: artifacts/mistral_full.plan.json)

    Returns:
        Dict with:
          - "tok_embeddings": descriptor for token embeddings (output.weight transposed)
          - "lm_head": descriptor for LM head (output.weight)

    Raises:
        RuntimeError: If plan file missing or output.weight not found
    """
    if plan_path is None:
        # Default to mistral_full.plan.json in helix-cdc/artifacts
        base_path = Path(__file__).parent.parent
        plan_path = base_path / "artifacts" / "mistral_full.plan.json"

    if not plan_path.exists():
        raise RuntimeError(f"Plan file not found: {plan_path}")

    plan = json.loads(plan_path.read_text())

    # Find output.weight (LM head)
    output_tensor = None
    for entry in plan:
        if entry.get("name") == "output.weight" and entry.get("role") == "output":
            output_tensor = entry
            break

    if output_tensor is None:
        raise RuntimeError("output.weight tensor not found in plan file")

    # Mistral uses tied embeddings: embedding matrix = output.weight.T
    shape = output_tensor["shape"]  # [hidden_size, vocab_size]
    hidden_size, vocab_size = shape

    # Note: We don't have "offset" in the plan, so set to 0 as placeholder
    # The actual regeneration will use the name + metadata from the plan
    return {
        "tok_embeddings": {
            "name": "output.weight",  # Same tensor, used transposed
            "shape": [vocab_size, hidden_size],  # Transposed: [32000, 4096]
            "size": output_tensor["size"],
            "offset": 0,  # Placeholder - regrow uses name not offset
            "transpose": True,  # Flag to indicate this needs transposing
            "qtype": output_tensor.get("qtype"),  # Quantization type (e.g., Q6_K)
            "dtype": output_tensor.get("dtype"),  # Data type (e.g., "float")
        },
        "lm_head": {
            "name": "output.weight",
            "shape": shape,  # [4096, 32000]
            "size": output_tensor["size"],
            "offset": 0,  # Placeholder
            "transpose": False,
            "qtype": output_tensor.get("qtype"),  # Quantization type (e.g., Q6_K)
            "dtype": output_tensor.get("dtype"),  # Data type (e.g., "float")
        }
    }


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


def describe_superglyph_substrate() -> SuperglyphSubstrateDescriptor:
    """
    Main entry point: describe the Superglyph substrate as visible to Helix-CDC.

    This is what other modules (substrate_inspector, dashboards, etc.) should call.
    """
    paths = SuperglyphPaths.detect()
    caps = _detect_capabilities(paths)
    genome = _detect_genome(paths.genome_root)

    return SuperglyphSubstrateDescriptor(
        name="Superglyph-Regenerative-Substrate",
        version="v1-arch-injected",
        paths=paths,
        capabilities=caps,
        genome=genome,
    )


def dump_superglyph_summary(pretty: bool = True) -> str:
    """
    Convenience helper to stringify the current Superglyph substrate descriptor.
    """
    desc = describe_superglyph_substrate()
    data = desc.to_dict()
    if pretty:
        return json.dumps(data, indent=2, sort_keys=True)
    return json.dumps(data, separators=(",", ":"))


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _main(argv: Optional[List[str]] = None) -> int:
    import argparse

    parser = argparse.ArgumentParser(
        prog="superglyph_bridge",
        description="Describe the Superglyph substrate as seen from Helix-CDC.",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Emit machine-readable JSON instead of human text.",
    )
    args = parser.parse_args(argv)

    if args.json:
        print(dump_superglyph_summary(pretty=False))
    else:
        desc = describe_superglyph_substrate()
        data = desc.to_dict()
        print("=== Superglyph Substrate Descriptor ===")
        print(f"name:    {data['name']}")
        print(f"version: {data['version']}")
        print()
        print("[paths]")
        for k, v in data["paths"].items():
            print(f"  {k}: {v}")
        print()
        print("[capabilities]")
        for k, v in data["capabilities"].items():
            print(f"  {k}: {v}")
        print()
        print("[genome]")
        for k, v in data["genome"].items():
            print(f"  {k}: {v}")

    return 0


if __name__ == "__main__":
    raise SystemExit(_main())
