#!/usr/bin/env python3
"""
SuperglyphAgenticCell - Autonomous agent with regenerative model capability

Architecture:
  - Genome: Superglyph (.hxz) with multiple model seeds
  - Phenotype: regrow() expresses needed model on-demand
  - Autonomy: Cell decides which capability to use per task

Integration:
  - Uses existing regrow() from helix_cdc.block_api
  - Compatible with AISwarmController
  - Receipt-driven (every operation logged)
"""
import sys, json, time, lzma, tarfile, hashlib, subprocess
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple

# Import Se-based phenotype router (replaces hardcoded routing)
try:
    from helix_cdc.agentic.se_router import PhenotypeRouter, route_task
    SE_ROUTER_AVAILABLE = True
except ImportError:
    SE_ROUTER_AVAILABLE = False

class SuperglyphAgenticCell:
    """
    Agentic cell with dynamic phenotype expression

    Example:
        cell = SuperglyphAgenticCell("~/archive/fused_capsules/echo_superglyph_@G13.hxz")
        phenotype = cell.express_phenotype("WizardCoder")
        result = cell.autonomous_task("code", "Generate sorting algorithm")
    """

    def __init__(self, genome_path: str):
        """
        Initialize cell with superglyph genome

        Args:
            genome_path: Path to superglyph .hxz file
        """
        self.genome_path = Path(genome_path).expanduser()
        if not self.genome_path.exists():
            raise FileNotFoundError(f"Genome not found: {self.genome_path}")

        self.genome_id = self._compute_genome_id()
        self.capabilities = self._scan_capabilities()
        self.active_phenotype = None
        self.expression_cache = {}

    def _compute_genome_id(self) -> str:
        """Fast genome fingerprint (header-only)"""
        with lzma.open(self.genome_path, "rb") as xf:
            header = xf.read(1024)
        return hashlib.sha256(header).hexdigest()[:16]

    def _scan_capabilities(self) -> Dict[str, Any]:
        """
        Scan genome for available model seeds (fast, no extraction)

        Returns:
            Dict with genome_size_mb, model_count, model_seeds list
        """
        lobes = []
        with lzma.open(self.genome_path, "rb") as xf:
            tf = tarfile.open(fileobj=xf, mode="r|*")
            for m in tf:
                if m.name.endswith(".hxz") and "model_lobes" in m.name:
                    lobes.append({
                        "name": Path(m.name).name,
                        "size": m.size,
                        "path": m.name
                    })

        return {
            "genome_size_mb": self.genome_path.stat().st_size / 1024 / 1024,
            "model_count": len(lobes),
            "model_seeds": lobes
        }

    def _extract_seed(self, seed_name: str) -> bytes:
        """
        Extract specific seed from genome

        Args:
            seed_name: Seed filename or pattern to match

        Returns:
            Seed bytes
        """
        with lzma.open(self.genome_path, "rb") as xf:
            tf = tarfile.open(fileobj=xf, mode="r|*")
            for m in tf:
                if seed_name in m.name and m.name.endswith(".hxz"):
                    return tf.extractfile(m).read()
        raise FileNotFoundError(f"Seed {seed_name} not found in genome")

    def express_phenotype(self, model_pattern: str) -> Dict[str, Any]:
        """
        Express a phenotype (regenerate model from seed)

        Args:
            model_pattern: Pattern to match seed name (e.g., "Mistral", "WizardCoder")

        Returns:
            Phenotype metadata (seed extracted, ready for regrow)
        """
        # Check cache first
        if model_pattern in self.expression_cache:
            return self.expression_cache[model_pattern]

        # Find matching seed
        matched = None
        for seed in self.capabilities["model_seeds"]:
            if model_pattern.lower() in seed["name"].lower():
                matched = seed
                break

        if not matched:
            raise ValueError(f"No seed matches '{model_pattern}' in genome")

        # Extract seed
        t0 = time.time()
        seed_bytes = self._extract_seed(matched["name"])
        extract_time = (time.time() - t0) * 1000

        # Create phenotype metadata
        phenotype = {
            "model_name": matched["name"],
            "seed_size_bytes": len(seed_bytes),
            "seed_sha256": hashlib.sha256(seed_bytes).hexdigest(),
            "expressed_at": time.time(),
            "extract_time_ms": extract_time,
            "status": "seed_extracted",
            "genome_id": self.genome_id
        }

        # Cache it
        self.expression_cache[model_pattern] = phenotype
        self.active_phenotype = model_pattern

        return phenotype

    def regrow(self, model_pattern: str) -> bytes:
        """
        Full regeneration: extract seed + regrow to model weights

        Args:
            model_pattern: Model name pattern

        Returns:
            Regenerated model weights (bytes)
        """
        phenotype = self.express_phenotype(model_pattern)
        seed_bytes = self._extract_seed(phenotype["model_name"])

        # Use regrow_cli.sh (already wired to regrow_q4km)
        regrow_script = Path(__file__).parent.parent.parent / "superglyph_lab" / "tools" / "regrow_cli.sh"
        if not regrow_script.exists():
            # Fallback: just return seed (placeholder for demo)
            return seed_bytes

        # Call regrow
        result = subprocess.run(
            [str(regrow_script)],
            input=seed_bytes,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )

        if result.returncode != 0:
            raise RuntimeError(f"Regrow failed: {result.stderr.decode()}")

        return result.stdout

    def regrow_and_test(self, phenotype_name: str, output_dir: Optional[str] = None) -> Dict[str, Any]:
        """
        Full regen + test with KAT receipt (WO:FIX-03 integration)

        This method:
        1. Extracts seed from genome
        2. Calls regrow_cli.sh to regenerate model
        3. Calls regrow_kat.py to produce KAT receipt
        4. Returns receipt with verification results

        Args:
            phenotype_name: Model name pattern (e.g., "Mistral", "WizardCoder")
            output_dir: Where to write regrown model + receipt (default: /tmp/regrown/)

        Returns:
            KAT receipt dict (with sha256_match, regen_time, etc.)
        """
        # 1. Extract seed
        phenotype = self.express_phenotype(phenotype_name)
        seed_bytes = self._extract_seed(phenotype["model_name"])

        # 2. Set up paths
        output_dir = Path(output_dir) if output_dir else Path("/tmp/regrown")
        output_dir.mkdir(parents=True, exist_ok=True)

        seed_path = output_dir / f"{phenotype_name}.hxz"
        model_path = output_dir / f"{phenotype_name}.gguf"
        receipt_path = output_dir / f"{phenotype_name}_kat.json"

        # Write seed
        with open(seed_path, 'wb') as f:
            f.write(seed_bytes)

        # 3. Call regrow_cli.sh
        regrow_script = Path(__file__).parent.parent.parent / "superglyph_lab" / "tools" / "regrow_cli.sh"
        if not regrow_script.exists():
            return {
                "status": "error",
                "error": f"regrow_cli.sh not found at {regrow_script}",
                "phenotype": phenotype_name
            }

        result = subprocess.run(
            [str(regrow_script), str(seed_path), str(model_path)],
            capture_output=True,
            text=True
        )

        if result.returncode != 0:
            return {
                "status": "error",
                "error": f"Regrow failed: {result.stderr}",
                "phenotype": phenotype_name
            }

        # 4. Call regrow_kat.py for verification
        kat_script = Path(__file__).parent.parent.parent / "superglyph_lab" / "tools" / "regrow_kat.py"
        if not kat_script.exists():
            return {
                "status": "warning",
                "warning": f"regrow_kat.py not found, skipping KAT verification",
                "phenotype": phenotype_name,
                "model_path": str(model_path)
            }

        kat_result = subprocess.run(
            ["python3", str(kat_script), str(model_path), str(receipt_path)],
            capture_output=True,
            text=True
        )

        if kat_result.returncode != 0:
            return {
                "status": "error",
                "error": f"KAT verification failed: {kat_result.stderr}",
                "phenotype": phenotype_name,
                "model_path": str(model_path)
            }

        # 5. Load and return receipt
        with open(receipt_path) as f:
            receipt = json.load(f)

        return receipt

    def autonomous_task(self, task_type: str, task_data: Any) -> Dict[str, Any]:
        """
        Execute autonomous task with appropriate phenotype

        Uses Se-based learned router (90.7% accuracy) instead of hardcoded if/else.
        Falls back to heuristic routing if router unavailable.

        Args:
            task_type: "code", "instruct", "reasoning", etc.
            task_data: Task-specific data

        Returns:
            Task result with receipt
        """
        t0 = time.time()

        # Build available phenotypes dict
        available_phenotypes = {
            seed["name"]: seed for seed in self.capabilities["model_seeds"]
        }

        routing_metadata = {}

        # Use Se-based learned router if available
        if SE_ROUTER_AVAILABLE and available_phenotypes:
            phenotype_name, routing_metadata = route_task(
                task_type=task_type,
                task_data=task_data,
                available_phenotypes=available_phenotypes
            )
            # Find matching seed
            matched_seed = None
            for seed in self.capabilities["model_seeds"]:
                if phenotype_name.lower() in seed["name"].lower() or seed["name"].lower() in phenotype_name.lower():
                    matched_seed = seed
                    break
            if matched_seed:
                phenotype = self.express_phenotype(matched_seed["name"][:20])
            elif self.capabilities["model_seeds"]:
                phenotype = self.express_phenotype(self.capabilities["model_seeds"][0]["name"][:10])
            else:
                raise ValueError("No model seeds available in genome")
        else:
            # Fallback: hardcoded routing (legacy, 73.6% accuracy)
            if task_type == "code":
                phenotype = self.express_phenotype("WizardCoder")
            elif task_type in ["instruct", "reasoning"]:
                phenotype = self.express_phenotype("mistral")
            else:
                if self.capabilities["model_seeds"]:
                    phenotype = self.express_phenotype(self.capabilities["model_seeds"][0]["name"][:10])
                else:
                    raise ValueError("No model seeds available in genome")
            routing_metadata = {"routing_method": "hardcoded_fallback"}

        execution_time = (time.time() - t0) * 1000

        # Return receipt with routing metadata
        receipt = {
            "receipt_type": "AUTONOMOUS_TASK",
            "task_type": task_type,
            "task_data": str(task_data)[:100],  # Truncate for receipt
            "phenotype_used": phenotype["model_name"],
            "seed_size_bytes": phenotype["seed_size_bytes"],
            "seed_sha256": phenotype["seed_sha256"],
            "genome_id": self.genome_id,
            "execution_time_ms": execution_time,
            "timestamp": time.time(),
        }

        # Include Se-router metadata if available
        if routing_metadata:
            receipt["routing"] = routing_metadata

        return receipt

    def replicate(self) -> 'SuperglyphAgenticCell':
        """
        Fast cell replication (reference genome, don't copy)

        Returns:
            New cell instance (shares same genome)
        """
        return SuperglyphAgenticCell(str(self.genome_path))

    def get_receipt(self) -> Dict[str, Any]:
        """
        Get cell receipt (genome info + capabilities)

        Returns:
            Cell receipt
        """
        return {
            "receipt_type": "SUPERGLYPH_AGENTIC_CELL",
            "genome_id": self.genome_id,
            "genome_path": str(self.genome_path),
            "genome_size_mb": self.capabilities["genome_size_mb"],
            "model_count": self.capabilities["model_count"],
            "model_seeds": [s["name"] for s in self.capabilities["model_seeds"]],
            "active_phenotype": self.active_phenotype,
            "timestamp": time.time()
        }


def demo():
    """Quick demo of SuperglyphAgenticCell"""
    import os

    genome = Path.home() / "archive/fused_capsules/echo_superglyph_@G13_77dedf82.hxz"

    if not genome.exists():
        print(f"ERROR: Genome not found: {genome}")
        return 1

    print("=" * 70)
    print("  🧬 SUPERGLYPH AGENTIC CELL DEMO")
    print("=" * 70)

    # Create cell
    print(f"\n📦 Initializing cell with genome: {genome.name}")
    cell = SuperglyphAgenticCell(str(genome))

    # Show capabilities
    receipt = cell.get_receipt()
    print(f"\n✅ Cell initialized:")
    print(f"   Genome ID: {receipt['genome_id']}")
    print(f"   Size: {receipt['genome_size_mb']:.1f}MB")
    print(f"   Models: {receipt['model_count']}")
    print(f"   Capabilities: {', '.join(receipt['model_seeds'][:3])}... (+{receipt['model_count']-3} more)")

    # Express phenotype
    print(f"\n🧬 Expressing phenotype: WizardCoder")
    phenotype = cell.express_phenotype("WizardCoder")
    print(f"   ✅ Seed extracted: {phenotype['seed_size_bytes']} bytes")
    print(f"   ⏱️  Extract time: {phenotype['extract_time_ms']:.2f}ms")
    print(f"   🔑 SHA256: {phenotype['seed_sha256'][:16]}...")

    # Autonomous task
    print(f"\n🤖 Executing autonomous task (code generation)")
    result = cell.autonomous_task("code", "Generate Python sorting algorithm")
    print(f"   ✅ Task completed")
    print(f"   Used: {result['phenotype_used']}")
    print(f"   Time: {result['execution_time_ms']:.2f}ms")

    # Replication
    print(f"\n🔬 Cell replication test")
    cell2 = cell.replicate()
    print(f"   ✅ New cell created (shares genome)")
    print(f"   Genome ID: {cell2.genome_id}")

    print("\n" + "=" * 70)
    print("  ✅ DEMO COMPLETE")
    print("=" * 70)

    return 0


if __name__ == "__main__":
    exit(demo())
