#!/usr/bin/env python3
"""
mistral_backends.py

WO-MISTRAL-REGEN-BEHAVIOR-FIBPI-SE-01:
Backend implementations for Mistral models that implement the BackendModel
protocol for use with HelixOrganism.

Two backends:
- GGUFMistralBackend: Wraps llama-cli for CPU inference of GGUF models
- RegenMistralBackend: Same, but specifically tracks regen metadata

Usage:
    from helix_cdc.loaders.mistral_backends import GGUFMistralBackend
    from helix_cdc.agentic.organism_adapter import HelixOrganism

    backend = GGUFMistralBackend.from_file(
        name="vendor_mistral",
        gguf_path="/path/to/mistral.gguf"
    )
    organism = HelixOrganism(backend=backend, organism_id="mistral_v01")
    result = organism.run_trial(
        episode_id="test",
        prompt_id="q1",
        prompt_text="What is 2+2?",
        se_mode="hot"
    )
"""

from __future__ import annotations

import hashlib
import json
import os
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional

# Default llama-cli path
LLAMA_CLI_PATH = os.environ.get(
    "LLAMA_CLI_PATH",
    "/home/voidstr3m33/llama.cpp/build/bin/llama-cli"
)


@dataclass
class GGUFMistralBackend:
    """
    GGUF Mistral backend using llama-cli for inference.

    Implements the BackendModel protocol from organism_adapter.py:
    - name: str
    - generate(prompt, max_tokens, temperature) -> str
    """
    name: str
    gguf_path: Path
    context_size: int = 4096
    threads: int = 4
    metadata: Dict[str, Any] = field(default_factory=dict)

    # Internal state
    _sha256_prefix: str = ""
    _size_bytes: int = 0
    _verified: bool = False

    @classmethod
    def from_file(
        cls,
        name: str,
        gguf_path: str,
        context_size: int = 4096,
        threads: int = 4,
    ) -> "GGUFMistralBackend":
        """
        Create backend from GGUF file path.

        Args:
            name: Backend identifier (e.g., "vendor_mistral", "regen_mistral")
            gguf_path: Path to GGUF model file
            context_size: Context window size
            threads: Number of threads for inference

        Returns:
            GGUFMistralBackend ready for generate()
        """
        path = Path(gguf_path)
        if not path.exists():
            raise FileNotFoundError(f"GGUF file not found: {path}")

        # Compute SHA256 prefix (first 1MB)
        with open(path, "rb") as f:
            sha256_prefix = hashlib.sha256(f.read(1024 * 1024)).hexdigest()[:16]

        size_bytes = path.stat().st_size

        backend = cls(
            name=name,
            gguf_path=path,
            context_size=context_size,
            threads=threads,
            metadata={
                "sha256_prefix": sha256_prefix,
                "size_bytes": size_bytes,
                "size_gb": size_bytes / 1e9,
            },
        )
        backend._sha256_prefix = sha256_prefix
        backend._size_bytes = size_bytes
        backend._verified = True

        return backend

    def generate(
        self,
        prompt: str,
        max_tokens: int = 128,
        temperature: float = 0.0,
    ) -> str:
        """
        Generate completion using llama-cli.

        Args:
            prompt: Input text
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature (0.0 = greedy)

        Returns:
            Generated text
        """
        if not Path(LLAMA_CLI_PATH).exists():
            raise FileNotFoundError(
                f"llama-cli not found at {LLAMA_CLI_PATH}\n"
                "Set LLAMA_CLI_PATH environment variable to correct path"
            )

        # Write prompt to temp file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            f.write(prompt)
            prompt_file = f.name

        try:
            cmd = [
                LLAMA_CLI_PATH,
                "-m", str(self.gguf_path),
                "-f", prompt_file,
                "-n", str(max_tokens),
                "--temp", str(temperature),
                "-c", str(self.context_size),
                "-t", str(self.threads),
                "--no-display-prompt",
            ]

            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=120,
            )

            if result.returncode != 0:
                error_msg = result.stderr[:500] if result.stderr else "Unknown error"
                raise RuntimeError(f"llama-cli error: {error_msg}")

            return result.stdout.strip()

        finally:
            os.unlink(prompt_file)

    def get_info(self) -> Dict[str, Any]:
        """Get backend metadata."""
        return {
            "name": self.name,
            "gguf_path": str(self.gguf_path),
            "sha256_prefix": self._sha256_prefix,
            "size_bytes": self._size_bytes,
            "context_size": self.context_size,
            "threads": self.threads,
            "verified": self._verified,
            **self.metadata,
        }


@dataclass
class RegenMistralBackend(GGUFMistralBackend):
    """
    Regenerated Mistral backend.

    Same as GGUFMistralBackend but tracks:
    - regen_source: Where the regen came from (e.g., "seed_487b")
    - regen_schema: Schema version used for regeneration
    - is_regen: True (for distinguishing from vendor)
    """
    regen_source: str = ""
    regen_schema: str = ""

    @classmethod
    def from_file(
        cls,
        name: str,
        gguf_path: str,
        regen_source: str = "unknown",
        regen_schema: str = "v1",
        context_size: int = 4096,
        threads: int = 4,
    ) -> "RegenMistralBackend":
        """Create regen backend from GGUF file."""
        path = Path(gguf_path)
        if not path.exists():
            raise FileNotFoundError(f"GGUF file not found: {path}")

        with open(path, "rb") as f:
            sha256_prefix = hashlib.sha256(f.read(1024 * 1024)).hexdigest()[:16]

        size_bytes = path.stat().st_size

        backend = cls(
            name=name,
            gguf_path=path,
            context_size=context_size,
            threads=threads,
            regen_source=regen_source,
            regen_schema=regen_schema,
            metadata={
                "sha256_prefix": sha256_prefix,
                "size_bytes": size_bytes,
                "size_gb": size_bytes / 1e9,
                "is_regen": True,
                "regen_source": regen_source,
                "regen_schema": regen_schema,
            },
        )
        backend._sha256_prefix = sha256_prefix
        backend._size_bytes = size_bytes
        backend._verified = True

        return backend


# =============================================================================
# Convenience loaders for known model paths
# =============================================================================

def load_vendor_mistral(
    path: Optional[str] = None,
    name: str = "vendor_mistral",
) -> GGUFMistralBackend:
    """
    Load vendor Mistral-7B from default or specified path.

    Default path: artifacts/capsules/mistral_FROM_CLONE.gguf
    """
    if path is None:
        # Try default paths
        candidates = [
            Path("/home/voidstr3m33/helix-cdc/artifacts/capsules/mistral_FROM_CLONE.gguf"),
            Path("/home/voidstr3m33/mistral_canonical_decoded.gguf"),
            Path("artifacts/capsules/mistral_FROM_CLONE.gguf"),
        ]
        for p in candidates:
            if p.exists():
                path = str(p)
                break

        if path is None:
            raise FileNotFoundError(
                "Vendor Mistral GGUF not found at default paths:\n" +
                "\n".join(f"  - {p}" for p in candidates)
            )

    return GGUFMistralBackend.from_file(name=name, gguf_path=path)


def load_regen_mistral(
    path: Optional[str] = None,
    name: str = "regen_mistral",
    regen_source: str = "seed_487b",
) -> RegenMistralBackend:
    """
    Load regenerated Mistral from default or specified path.

    Default path: artifacts/mistral_from_seed_cdna_v1_regen_cli.gguf

    NOTE: The seed3 procedural artifact has been quarantined. Do not
    use mistral_regen_seed3_v01/mistral_seed3_v01.gguf.
    """
    if path is None:
        candidates = [
            # Current verified regen from Helix64 manifest
            Path("/home/voidstr3m33/helix-cdc/artifacts/mistral_from_seed_cdna_v1_regen_cli.gguf"),
            Path("artifacts/mistral_from_seed_cdna_v1_regen_cli.gguf"),
            # Fallback to historical clone (not procedural, but not preferred)
            Path("/home/voidstr3m33/helix-cdc/artifacts/capsules/mistral_regen_FROM_CLONE.gguf"),
        ]
        for p in candidates:
            if p.exists():
                path = str(p)
                break

        if path is None:
            raise FileNotFoundError(
                "Regen Mistral GGUF not found at default paths:\n" +
                "\n".join(f"  - {p}" for p in candidates)
            )

    return RegenMistralBackend.from_file(
        name=name,
        gguf_path=path,
        regen_source=regen_source,
        regen_schema="mistral_regen_from_seed3.v03",
    )


# =============================================================================
# Module test
# =============================================================================

def _test_backends():
    """Test backend loading and basic generation."""
    print("Testing Mistral backends...")

    # Test vendor backend
    try:
        vendor = load_vendor_mistral()
        print(f"[OK] Loaded vendor backend: {vendor.name}")
        print(f"     Path: {vendor.gguf_path}")
        print(f"     Size: {vendor._size_bytes / 1e9:.2f} GB")
        print(f"     SHA256: {vendor._sha256_prefix}...")
    except FileNotFoundError as e:
        print(f"[SKIP] Vendor not found: {e}")
        vendor = None

    # Test regen backend
    try:
        regen = load_regen_mistral()
        print(f"[OK] Loaded regen backend: {regen.name}")
        print(f"     Path: {regen.gguf_path}")
        print(f"     Size: {regen._size_bytes / 1e9:.2f} GB")
        print(f"     Source: {regen.regen_source}")
    except FileNotFoundError as e:
        print(f"[SKIP] Regen not found: {e}")
        regen = None

    # Quick generation test (if llama-cli available)
    if vendor and Path(LLAMA_CLI_PATH).exists():
        print("\nTesting generation (vendor)...")
        try:
            result = vendor.generate("Hello,", max_tokens=8, temperature=0.0)
            print(f"  Prompt: 'Hello,'")
            print(f"  Output: '{result[:50]}...'")
            print("[OK] Generation works")
        except Exception as e:
            print(f"[FAIL] Generation error: {e}")
    else:
        print("\n[SKIP] Generation test (llama-cli not available or vendor not loaded)")

    print("\nPASS: Backend loading test")


if __name__ == "__main__":
    _test_backends()
