"""
Edge Inference Entry Point (WO-EDGE-DEPTH-ROUTER-01)

Unified entry point for edge device inference. Routes queries based on
Symbolic Entropy (Se) to appropriate depth, then executes either local
streaming inference or cloud fallback.

Architecture:
    1. Compute/estimate Se from input
    2. Route Se to depth (with RAM constraints)
    3. Execute local multi-block forward OR cloud fallback
    4. Return result with full receipt

Usage:
    from helix_cdc.regrow.edge_inference import edge_inference

    result = edge_inference(
        x=input_activations,
        prompt="What is quantum entanglement?",
        manifest_path="seeds/hybrid_manifest_v2.json",
        available_ram_mb=256,
    )
"""

import os
import time
import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Dict, Any, List, Union, Tuple
from datetime import datetime, timezone

import numpy as np

from .depth_router import (
    InferenceDepth,
    DepthRouteResult,
    route_se_to_depth,
    get_block_indices_for_depth,
    estimate_se_from_prompt,
)
from .gemini_fallback import (
    GeminiFallbackResult,
    query_gemini,
    check_gemini_available,
)


# =============================================================================
# EdgeInferenceResult Dataclass
# =============================================================================

@dataclass
class EdgeInferenceResult:
    """
    Result from edge inference pipeline.

    Contains the inference result, routing decision, and comprehensive
    receipt for claim hygiene.
    """
    # Inference output
    output: Optional[np.ndarray] = None  # Local inference output
    response_text: Optional[str] = None  # Cloud fallback response
    success: bool = False

    # Routing decision
    route_result: Optional[DepthRouteResult] = None
    inference_mode: str = ""  # "local" or "cloud"

    # For local inference
    blocks_executed: List[int] = field(default_factory=list)
    block_receipts: List[Dict[str, Any]] = field(default_factory=list)

    # For cloud fallback
    gemini_result: Optional[GeminiFallbackResult] = None

    # Input metadata
    prompt: str = ""
    prompt_hash: str = ""
    se_score: float = 0.0
    se_source: str = ""  # "computed" or "estimated"

    # Resource usage
    available_ram_mb: float = 0.0
    estimated_ram_used_mb: float = 0.0

    # Timing
    total_time_ms: float = 0.0
    routing_time_ms: float = 0.0
    inference_time_ms: float = 0.0

    # Error handling
    error_message: Optional[str] = None

    # Receipt schema
    schema: str = "edge_inference_result_v1"
    work_order: str = "WO-EDGE-DEPTH-ROUTER-01"
    timestamp: str = ""

    # Chat generation fields (WO-TOKENIZER-BRIDGE-01)
    generated_tokens: List[int] = field(default_factory=list)
    generated_text: str = ""
    tokenizer_receipts: List[Dict[str, Any]] = field(default_factory=list)
    per_step_logits_summary: List[Dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        result = {
            "schema": self.schema,
            "work_order": self.work_order,
            "timestamp": self.timestamp,
            "success": self.success,
            "inference_mode": self.inference_mode,
            "prompt_hash": self.prompt_hash,
            "se_score": self.se_score,
            "se_source": self.se_source,
            "routing": self.route_result.to_dict() if self.route_result else None,
            "available_ram_mb": self.available_ram_mb,
            "estimated_ram_used_mb": self.estimated_ram_used_mb,
            "timing": {
                "total_ms": self.total_time_ms,
                "routing_ms": self.routing_time_ms,
                "inference_ms": self.inference_time_ms,
            },
            "error_message": self.error_message,
        }

        if self.inference_mode == "local":
            result["local_inference"] = {
                "blocks_executed": self.blocks_executed,
                "output_shape": list(self.output.shape) if self.output is not None else None,
                "output_hash": _hash_array(self.output) if self.output is not None else None,
            }
        elif self.inference_mode == "cloud":
            result["cloud_fallback"] = self.gemini_result.to_dict() if self.gemini_result else None

        # Chat generation fields (WO-TOKENIZER-BRIDGE-01)
        if self.generated_tokens:
            result["chat_generation"] = {
                "generated_tokens": self.generated_tokens,
                "generated_text": self.generated_text,
                "token_count": len(self.generated_tokens),
                "tokenizer_receipts": self.tokenizer_receipts,
                "per_step_logits_summary": self.per_step_logits_summary,
            }

        return result

    def save_receipt(self, path: Union[str, Path]) -> None:
        """Save result as JSON receipt."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)


# =============================================================================
# Main Entry Point
# =============================================================================

def edge_inference(
    x: Optional[np.ndarray] = None,
    prompt: str = "",
    manifest_path: Optional[Union[str, Path]] = None,
    available_ram_mb: float = 512.0,
    force_local: bool = False,
    se_score: Optional[float] = None,
    gguf_path: Optional[Union[str, Path]] = None,
    d_head: int = 128,
    verify_policy: str = "trust_cached",
    dry_run: bool = False,
) -> EdgeInferenceResult:
    """
    Run edge inference with Se-based depth routing.

    This is the main entry point for edge device inference. It:
    1. Computes or estimates Se from the input
    2. Routes to appropriate depth based on Se and RAM
    3. Executes local streaming inference or cloud fallback
    4. Returns comprehensive result with receipt

    Args:
        x: Input activations for local inference [batch, seq, d_model].
            If None and local inference is needed, will create dummy input.
        prompt: Input prompt text. Used for Se estimation and cloud fallback.
        manifest_path: Path to hybrid_manifest_v2.json for local inference.
        available_ram_mb: Available RAM in megabytes for depth routing.
        force_local: If True, never use cloud fallback.
        se_score: Pre-computed Se score. If None, estimated from prompt.
        gguf_path: Path to GGUF file for norm weights (local inference).
        d_head: Head dimension for attention (default 128 for Mistral).
        verify_policy: Block verification policy for local inference.
        dry_run: If True, don't actually run inference (for testing routing).

    Returns:
        EdgeInferenceResult with output and full receipt.

    Examples:
        # Local inference with sufficient RAM
        result = edge_inference(
            x=input_tensor,
            prompt="What is 2+2?",
            manifest_path="seeds/hybrid_manifest_v2.json",
            available_ram_mb=256,
        )

        # Cloud fallback for complex query
        result = edge_inference(
            prompt="Explain the philosophical implications of quantum mechanics",
            available_ram_mb=64,
        )

        # Dry run to test routing
        result = edge_inference(
            prompt="Hello world",
            available_ram_mb=128,
            dry_run=True,
        )
    """
    start_time = time.perf_counter()
    timestamp = datetime.now(timezone.utc).isoformat()

    # Initialize result
    result = EdgeInferenceResult(
        prompt=prompt,
        prompt_hash=_hash_prompt(prompt),
        available_ram_mb=available_ram_mb,
        timestamp=timestamp,
    )

    try:
        # Step 1: Compute or estimate Se
        if se_score is not None:
            result.se_score = se_score
            result.se_source = "provided"
        else:
            result.se_score = estimate_se_from_prompt(prompt)
            result.se_source = "estimated_from_prompt"

        # Step 2: Route to depth
        routing_start = time.perf_counter()
        route_result = route_se_to_depth(
            se_score=result.se_score,
            available_ram_mb=available_ram_mb,
            force_local=force_local,
        )
        result.route_result = route_result
        result.routing_time_ms = (time.perf_counter() - routing_start) * 1000

        # Step 3: Execute inference based on routing
        inference_start = time.perf_counter()

        if route_result.depth == InferenceDepth.CLOUD:
            # Cloud fallback
            result.inference_mode = "cloud"

            if dry_run:
                result.success = True
                result.response_text = "[DRY RUN] Would query Gemini API"
            else:
                gemini_result = query_gemini(prompt)
                result.gemini_result = gemini_result
                result.success = gemini_result.success
                result.response_text = gemini_result.response_text
                if not gemini_result.success:
                    result.error_message = gemini_result.error_message

        else:
            # Local inference
            result.inference_mode = "local"
            result.blocks_executed = get_block_indices_for_depth(route_result.depth)
            result.estimated_ram_used_mb = route_result.estimated_ram_mb

            if dry_run:
                result.success = True
                result.output = None
            else:
                # Run local inference
                output, receipts = _run_local_inference(
                    x=x,
                    manifest_path=manifest_path,
                    block_indices=result.blocks_executed,
                    gguf_path=gguf_path,
                    d_head=d_head,
                    verify_policy=verify_policy,
                )
                result.output = output
                result.block_receipts = [r.to_dict() if hasattr(r, 'to_dict') else r for r in receipts]
                result.success = True

        result.inference_time_ms = (time.perf_counter() - inference_start) * 1000

    except Exception as e:
        result.success = False
        result.error_message = str(e)

    result.total_time_ms = (time.perf_counter() - start_time) * 1000
    return result


# =============================================================================
# Local Inference Helper
# =============================================================================

def _run_local_inference(
    x: Optional[np.ndarray],
    manifest_path: Optional[Union[str, Path]],
    block_indices: List[int],
    gguf_path: Optional[Union[str, Path]],
    d_head: int,
    verify_policy: str,
) -> Tuple[np.ndarray, List[Any]]:
    """
    Run local streaming inference through transformer blocks.

    Args:
        x: Input activations. If None, creates dummy input.
        manifest_path: Path to manifest.
        block_indices: Which blocks to run.
        gguf_path: Path to GGUF for norm weights.
        d_head: Head dimension.
        verify_policy: Verification policy.

    Returns:
        (output_tensor, list_of_receipts)
    """
    # Import here to avoid circular imports and allow graceful degradation
    try:
        from .stream_transformer_block import stream_multi_block_forward
    except ImportError as e:
        raise ImportError(
            f"stream_transformer_block not available: {e}. "
            "Local inference requires full helix_cdc installation."
        )

    # Validate manifest path
    if manifest_path is None:
        raise ValueError("manifest_path required for local inference")

    manifest_path = Path(manifest_path)
    if not manifest_path.exists():
        raise FileNotFoundError(f"Manifest not found: {manifest_path}")

    # Create dummy input if needed
    if x is None:
        # Mistral 7B: d_model=4096
        x = np.random.randn(1, 32, 4096).astype(np.float32)

    # Run multi-block forward
    output, receipts = stream_multi_block_forward(
        X=x,
        manifest_path=manifest_path,
        block_indices=block_indices,
        d_head=d_head,
        verify_policy=verify_policy,
        gguf_path=gguf_path,
    )

    return output, receipts


# =============================================================================
# Chat Inference (WO-TOKENIZER-BRIDGE-01)
# =============================================================================

def edge_inference_chat(
    prompt: str,
    gguf_path: Union[str, Path],
    manifest_path: Union[str, Path],
    max_tokens: int = 64,
    available_ram_mb: float = 512.0,
    temperature: float = 1.0,
    d_head: int = 128,
    verify_policy: str = "trust_cached",
) -> EdgeInferenceResult:
    """
    Text-to-text generation using native CDNA v2 streaming.

    This is the "front door" for native CDNA chat. Unlike edge_inference()
    which takes activations, this takes text and returns text.

    Pipeline:
      1. Tokenize prompt (MistralTokenizerBridge.encode)
      2. Embed tokens (MistralTokenizerBridge.lookup_embeddings)
      3. Route by Se (route_se_to_depth)
      4. stream_multi_block_forward(X, block_indices)
      5. Project to logits (MistralTokenizerBridge.project_to_logits)
      6. Greedy decode, append token
      7. Repeat until max_tokens or EOS
      8. Decode to text (MistralTokenizerBridge.decode)

    Args:
        prompt: Input text prompt.
        gguf_path: Path to GGUF file (for tokenizer and weights).
        manifest_path: Path to hybrid_manifest_v2.json (for CDNA streaming).
        max_tokens: Maximum tokens to generate (default 64).
        available_ram_mb: Available RAM for depth routing (default 512).
        temperature: Temperature for logits (default 1.0).
        d_head: Head dimension for attention (default 128 for Mistral).
        verify_policy: Block verification policy.

    Returns:
        EdgeInferenceResult with generated_text and full receipts.

    Example:
        result = edge_inference_chat(
            prompt="What is 2+2?",
            gguf_path="model.gguf",
            manifest_path="seeds/hybrid_manifest_v2.json",
            max_tokens=32,
        )
        print(result.generated_text)
    """
    start_time = time.perf_counter()
    timestamp = datetime.now(timezone.utc).isoformat()

    # Initialize result
    result = EdgeInferenceResult(
        prompt=prompt,
        prompt_hash=_hash_prompt(prompt),
        available_ram_mb=available_ram_mb,
        timestamp=timestamp,
        work_order="WO-TOKENIZER-BRIDGE-01",
        schema="edge_inference_chat_result_v1",
    )

    try:
        # Import tokenizer bridge
        from .tokenizer_bridge import MistralTokenizerBridge

        # Step 1: Get tokenizer bridge
        bridge = MistralTokenizerBridge.get_bridge(gguf_path)

        # Step 2: Encode prompt
        token_ids, encode_receipt = bridge.encode(prompt, add_bos=True)
        result.tokenizer_receipts.append(encode_receipt.to_dict())

        if encode_receipt.status != "PASS":
            result.error_message = f"Encode failed: {encode_receipt.error_message}"
            result.total_time_ms = (time.perf_counter() - start_time) * 1000
            return result

        # Step 3: Estimate Se and route
        routing_start = time.perf_counter()
        result.se_score = estimate_se_from_prompt(prompt)
        result.se_source = "estimated_from_prompt"

        route_result = route_se_to_depth(
            se_score=result.se_score,
            available_ram_mb=available_ram_mb,
            force_local=True,  # Chat always uses local
        )
        result.route_result = route_result
        result.routing_time_ms = (time.perf_counter() - routing_start) * 1000

        if route_result.depth == InferenceDepth.CLOUD:
            # Can't do chat with cloud fallback
            result.error_message = "Cloud fallback not supported for chat generation"
            result.total_time_ms = (time.perf_counter() - start_time) * 1000
            return result

        result.inference_mode = "local"
        result.blocks_executed = get_block_indices_for_depth(route_result.depth)
        result.estimated_ram_used_mb = route_result.estimated_ram_mb

        # Step 4: Get embeddings
        embeddings, embed_receipt = bridge.lookup_embeddings(token_ids)
        result.tokenizer_receipts.append(embed_receipt.to_dict())

        if embed_receipt.status != "PASS":
            result.error_message = f"Embedding lookup failed: {embed_receipt.error_message}"
            result.total_time_ms = (time.perf_counter() - start_time) * 1000
            return result

        # Get EOS token for stopping
        eos_token_id = bridge.get_eos_token_id()

        # Import stream_multi_block_forward
        from .stream_transformer_block import stream_multi_block_forward

        manifest_path = Path(manifest_path)
        gguf_path = Path(gguf_path)

        # Generation loop
        inference_start = time.perf_counter()
        current_ids = list(token_ids)
        generated_ids = []

        for step in range(max_tokens):
            t_step_start = time.perf_counter()

            # Get embeddings for current sequence
            t_embed_start = time.perf_counter()
            if step == 0:
                # Use cached embeddings from initial encode
                current_embeddings = embeddings
            else:
                # Re-embed full sequence (no KV cache yet)
                current_embeddings, _ = bridge.lookup_embeddings(current_ids)
            t_embed_ms = (time.perf_counter() - t_embed_start) * 1000

            # Add batch dimension if needed: [seq, d_model] -> [1, seq, d_model]
            if current_embeddings.ndim == 2:
                current_embeddings = current_embeddings.reshape(1, *current_embeddings.shape)

            # Run through transformer blocks
            t_forward_start = time.perf_counter()
            hidden, block_receipts = stream_multi_block_forward(
                X=current_embeddings,
                manifest_path=manifest_path,
                block_indices=result.blocks_executed,
                d_head=d_head,
                verify_policy=verify_policy,
                gguf_path=gguf_path,
            )
            t_forward_ms = (time.perf_counter() - t_forward_start) * 1000

            # Only keep block receipts from first step to avoid huge receipts
            if step == 0:
                result.block_receipts = [
                    r.to_dict() if hasattr(r, 'to_dict') else r
                    for r in block_receipts
                ]

            # Get last position hidden state
            if hidden.ndim == 3:
                last_hidden = hidden[0, -1, :]  # [d_model]
            else:
                last_hidden = hidden[-1, :]  # [d_model]

            # Project to logits
            t_lmhead_start = time.perf_counter()
            logits, logits_receipt = bridge.project_to_logits(
                last_hidden.reshape(1, -1),  # [1, d_model]
                temperature=temperature,
            )
            t_lmhead_ms = (time.perf_counter() - t_lmhead_start) * 1000

            t_step_ms = (time.perf_counter() - t_step_start) * 1000

            # Store logits summary with timing breakdown
            result.per_step_logits_summary.append({
                "step": step,
                "seq_len": len(current_ids),
                "top_token_id": logits_receipt.top_token_id,
                "top_token_logit": logits_receipt.top_token_logit,
                "embed_ms": round(t_embed_ms, 1),
                "forward_ms": round(t_forward_ms, 1),
                "lmhead_ms": round(t_lmhead_ms, 1),
                "step_ms": round(t_step_ms, 1),
            })

            if logits_receipt.status != "PASS":
                result.error_message = f"Logits projection failed at step {step}"
                break

            # Greedy decode: argmax
            next_token_id = logits_receipt.top_token_id

            # Append to sequence
            current_ids.append(next_token_id)
            generated_ids.append(next_token_id)

            # Check for EOS
            if next_token_id == eos_token_id:
                break

        result.inference_time_ms = (time.perf_counter() - inference_start) * 1000

        # Store generated tokens
        result.generated_tokens = generated_ids

        # Decode generated tokens to text
        if generated_ids:
            generated_text, decode_receipt = bridge.decode(generated_ids)
            result.generated_text = generated_text
            result.tokenizer_receipts.append(decode_receipt.to_dict())
        else:
            result.generated_text = ""

        result.success = True

    except Exception as e:
        result.success = False
        result.error_message = str(e)

    result.total_time_ms = (time.perf_counter() - start_time) * 1000
    return result


# =============================================================================
# Quick Inference Helpers
# =============================================================================

def quick_route(prompt: str, available_ram_mb: float = 512.0) -> Dict[str, Any]:
    """
    Quick routing check without running inference.

    Args:
        prompt: Input prompt text.
        available_ram_mb: Available RAM.

    Returns:
        Dict with routing decision summary.
    """
    se = estimate_se_from_prompt(prompt)
    route = route_se_to_depth(se, available_ram_mb)

    return {
        "prompt_preview": prompt[:50] + "..." if len(prompt) > 50 else prompt,
        "se_score": round(se, 3),
        "depth": route.depth.name,
        "blocks": route.block_count,
        "ram_mb": route.estimated_ram_mb,
        "was_downgraded": route.was_downgraded,
        "downgrade_reason": route.downgrade_reason,
    }


def check_edge_readiness() -> Dict[str, Any]:
    """
    Check if edge inference is ready.

    Returns:
        Dict with status of each component.
    """
    status = {
        "depth_router": True,  # Always available
        "gemini_fallback": False,
        "gemini_message": "",
        "local_inference": False,
        "local_message": "",
        "manifest_path": None,
    }

    # Check Gemini
    available, message = check_gemini_available()
    status["gemini_fallback"] = available
    status["gemini_message"] = message

    # Check local inference
    try:
        from .stream_transformer_block import stream_multi_block_forward
        status["local_inference"] = True
        status["local_message"] = "stream_multi_block_forward available"
    except ImportError as e:
        status["local_message"] = str(e)

    # Check for common manifest paths
    common_paths = [
        "seeds/hybrid_manifest_v2.json",
        "seeds/hybrid_manifest_v2_pilot.json",
    ]
    for p in common_paths:
        if Path(p).exists():
            status["manifest_path"] = p
            break

    return status


# =============================================================================
# Utilities
# =============================================================================

def _hash_prompt(prompt: str) -> str:
    """Generate SHA256 hash of prompt."""
    return hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:16]


def _hash_array(arr: np.ndarray) -> str:
    """Generate SHA256 hash of numpy array."""
    return hashlib.sha256(arr.tobytes()).hexdigest()[:16]
