from __future__ import annotations
from typing import Dict, List, Tuple, Any, Optional
import uuid
import time
import logging
from pathlib import Path

INTERFACE_VERSION = "1.0"

logger = logging.getLogger(__name__)


def load_gx(path: str) -> Tuple[dict, List[dict], bytes]:
    """Load a .gx file and return (manifest, segments, payload).

    Thin wrapper around runtime_executor.gx_loader.load_gx().
    Reconstructs segments list from manifest['codex_lineage']['segments'].

    Args:
        path: Path to .gx file

    Returns:
        Tuple of (manifest, segments_list, payload_bytes)

    Raises:
        FileNotFoundError: If .gx file doesn't exist
        RuntimeError: If manifest is malformed
    """
    from runtime_executor.gx_loader import load_gx as loader_load_gx

    manifest, payload = loader_load_gx(path)

    # Extract segments from manifest['codex_lineage']['segments']
    codex_lineage = manifest.get("codex_lineage", {})
    segments = codex_lineage.get("segments", [])

    if not segments:
        logger.warning(f"No segments found in {path}, continuing with empty list")

    return manifest, segments, payload


def normalize_segments(
    manifest: dict,
    raw_segments: List[dict],
    payload: bytes,
) -> List[dict]:
    """Normalize raw segments into canonical Segment schema.

    Converts raw segment metadata from codex_lineage.segments into normalized form
    with required keys: id, start_line, end_line, text, symbolic_lane, semantic_role.

    Assumptions:
    - Raw segments have: id, start, end (0-based line indices)
    - Optional: lane, role metadata
    - Text extraction is stubbed (reserved for future payload analysis)

    Args:
        manifest: GX manifest dict
        raw_segments: List of raw segment dicts from codex_lineage
        payload: Compressed payload bytes (reserved for text extraction)

    Returns:
        List of normalized segment dicts conforming to Segment schema
    """
    normalized = []

    for raw_seg in raw_segments:
        seg_id = str(raw_seg.get("id", "unknown"))
        start_line = int(raw_seg.get("start", 0))
        end_line = int(raw_seg.get("end", 0))

        # Explicit lane from metadata, or infer
        explicit_lane = raw_seg.get("lane")
        if explicit_lane is not None:
            symbolic_lane = int(explicit_lane)
        else:
            symbolic_lane = _infer_lane(raw_seg)

        # Clamp to valid range [0, 7]
        symbolic_lane = max(0, min(7, symbolic_lane))

        # Semantic role: infer from segment metadata
        semantic_role = _infer_semantic_role(raw_seg)

        # Text: stub for now; in production extract from payload via byte ranges
        text = f"[segment {seg_id}: lines {start_line}–{end_line}]"

        normalized_seg = {
            "id": seg_id,
            "start_line": start_line,
            "end_line": end_line,
            "text": text,
            "symbolic_lane": symbolic_lane,
            "semantic_role": semantic_role,
            "_raw": raw_seg,  # preserve original for debugging
        }

        normalized.append(normalized_seg)

    return normalized


def _infer_lane(segment: dict) -> int:
    """Infer lane assignment from segment metadata (heuristic).

    Rules (minimum viable):
    - Structural markers → lane 0
    - Main content → lane 1 (default)
    - Comments/annotations → lane 3
    - Hints → lane 4
    - Author/meta → lane 6
    - Time/epoch → lane 7

    Args:
        segment: Raw segment dict

    Returns:
        Lane id 0–7
    """
    seg_id = str(segment.get("id", "")).lower()

    # Structural
    if any(x in seg_id for x in ["struct", "class", "def", "header", "header"]):
        return 0

    # Annotations/comments
    if any(x in seg_id for x in ["comment", "annotation", "tag"]):
        return 3

    # Hints
    if any(x in seg_id for x in ["hint", "note", "tip", "warn"]):
        return 4

    # Author/meta
    if any(x in seg_id for x in ["meta", "author", "signature"]):
        return 6

    # Time/epoch
    if any(x in seg_id for x in ["epoch", "time", "date", "version"]):
        return 7

    # Default: semantic flow
    return 1


def _infer_semantic_role(segment: dict) -> str:
    """Infer semantic role from segment metadata (heuristic).

    Args:
        segment: Raw segment dict

    Returns:
        One of: "definition", "constraint", "example", "meta", "unknown"
    """
    seg_id = str(segment.get("id", "")).lower()

    if any(x in seg_id for x in ["def", "class", "function", "declaration"]):
        return "definition"

    if any(x in seg_id for x in ["constraint", "rule", "assertion", "require"]):
        return "constraint"

    if any(x in seg_id for x in ["example", "test", "sample", "demo"]):
        return "example"

    if any(x in seg_id for x in ["meta", "note", "comment", "annotation", "tag"]):
        return "meta"

    return "unknown"


def map_lanes(segments: List[dict]) -> Dict[int, List[dict]]:
    """Map normalized segments into 0–7 lane model.

    Organizes segments by symbolic_lane into a dict where:
    - Keys: lane numbers 0–7
    - Values: lists of segments assigned to that lane

    Lane semantics (from spec):
    - 0: structural_logic
    - 1: semantic_flow
    - 2: compression_residue
    - 3: symbolic_metadata
    - 4: execution_hints
    - 5: predictive_scaffolding
    - 6: contributor_imprint
    - 7: epoch_resonance

    Args:
        segments: List of normalized segment dicts

    Returns:
        Dict[int, List[dict]] where keys are 0–7
    """
    lanes: Dict[int, List[dict]] = {i: [] for i in range(8)}

    for seg in segments:
        lane = seg.get("symbolic_lane", 1)
        # Safety clamp
        lane = max(0, min(7, lane))
        lanes[lane].append(seg)

    return lanes


def build_envelope(
    manifest: dict,
    lanes: Dict[int, List[dict]],
    payload: bytes,
    context: Optional[dict] = None,
) -> dict:
    """Build ExecutionEnvelope for LAIN from components.

    Constructs the envelope that LAIN will consume. The envelope is immutable
    from LAIN's perspective and includes all necessary context.

    Args:
        manifest: GX manifest dict
        lanes: Dict[int, List[dict]] where keys are lane ids 0–7
        payload: Compressed payload bytes
        context: Optional context overrides (runtime_flags, epoch, cognitive_mode, invocation_id)

    Returns:
        ExecutionEnvelope dict ready for LAIN.execute()
    """
    if context is None:
        context = {}

    base_context = {
        "runtime_flags": context.get("runtime_flags", {}),
        "contributor": manifest.get("contributor", "unknown"),
        "epoch": context.get("epoch"),
        "cognitive_mode": context.get("cognitive_mode", "analyze"),
        "invocation_id": context.get("invocation_id", str(uuid.uuid4())),
        "interface_version": INTERFACE_VERSION,
    }

    envelope = {
        "manifest": manifest,
        "lanes": lanes,
        "payload": payload,
        "context": base_context,
    }

    return envelope


def execute_with_lain(envelope: dict) -> dict:
    """Execute ExecutionEnvelope through (stub) LAIN engine.

    This is a stub implementation that simulates LAIN cognition.
    In production, this would call the real LAIN runtime.

    Contract:
    - Does not mutate input envelope
    - Deterministic for a given envelope
    - Errors are returned in result['diagnostics']['errors'], not raised

    Args:
        envelope: ExecutionEnvelope dict from build_envelope()

    Returns:
        ExecutionResult dict with cognition_trace, fused_symbol, output_text, diagnostics
    """
    start_time = time.time()

    manifest = envelope.get("manifest", {})
    lanes = envelope.get("lanes", {})
    payload = envelope.get("payload", b"")
    context = envelope.get("context", {})

    # Initialize diagnostics
    lane_timings: Dict[int, float] = {}
    errors: List[dict] = []

    # Stub: simulate processing each lane
    for lane_id in sorted(lanes.keys()):
        lane_timings[lane_id] = 0.0

    # Build cognition trace (stub)
    cognition_trace = []

    # Step 0: Load
    cognition_trace.append({
        "step": 0,
        "lane": -1,
        "segment_id": None,
        "operation": "load_envelope",
        "input": {
            "lanes": sorted(lanes.keys()),
            "num_segments": sum(len(segs) for segs in lanes.values()),
            "manifest_version": manifest.get("version"),
        },
        "output": {},
        "note": "Loaded ExecutionEnvelope into LAIN stub.",
    })

    # Step 1: Process lanes
    num_segments = sum(len(segs) for segs in lanes.values())
    cognition_trace.append({
        "step": 1,
        "lane": -1,
        "segment_id": None,
        "operation": "process_lanes",
        "input": {"lanes": sorted(lanes.keys())},
        "output": {"processed_segments": num_segments},
        "note": f"Stub processed {num_segments} segments.",
    })

    # Synthesize fused_symbol from lanes and segments
    all_segments = []
    for lane_id in sorted(lanes.keys()):
        all_segments.extend(lanes[lane_id])

    key_points = [seg["id"] for seg in all_segments[:3]]

    fused_symbol = {
        "summary": f"GX→LAIN stub: {len(all_segments)} segments, {len(lanes)} lanes",
        "key_points": key_points,
        "constraints": [],
        "open_questions": ["Real LAIN cognition not yet implemented"],
    }

    # Output text
    contributor = manifest.get("contributor", "unknown")
    source = manifest.get("source_file", "unknown")
    output_text = (
        f"GX→LAIN Runtime Stub v{INTERFACE_VERSION}\n"
        f"Source: {source}\n"
        f"Contributor: {contributor}\n"
        f"Segments: {len(all_segments)}\n"
        f"Lanes: {len(lanes)}\n"
        f"Status: Stub execution (replace with real LAIN engine)\n"
    )

    elapsed = time.time() - start_time

    # Build diagnostics
    diagnostics = {
        "lane_timings": lane_timings,
        "errors": errors,
        "resonance": {},
        "interface_version": INTERFACE_VERSION,
        "elapsed": elapsed,
    }

    # Return ExecutionResult
    result = {
        "cognition_trace": cognition_trace,
        "fused_symbol": fused_symbol,
        "output_text": output_text,
        "diagnostics": diagnostics,
    }

    return result


def execute_gx_path(
    gx_path: str,
    context: Optional[dict] = None,
) -> dict:
    """Main entry point: load .gx file and execute through LAIN.

    Pipeline: load → normalize → map_lanes → build_envelope → execute

    Args:
        gx_path: Path to .gx file
        context: Optional context overrides

    Returns:
        ExecutionResult dict

    Raises:
        Exceptions from load_gx() if file is invalid or missing
    """
    # Load
    manifest, raw_segments, payload = load_gx(gx_path)

    # Normalize
    segments = normalize_segments(manifest, raw_segments, payload)

    # Map to lanes
    lanes = map_lanes(segments)

    # Build envelope
    envelope = build_envelope(manifest, lanes, payload, context)

    # Execute
    result = execute_with_lain(envelope)

    return result


def make_error(
    error_type: str,
    message: str,
    segment_id: Optional[str] = None,
    lane: Optional[int] = None,
    recoverable: bool = True,
) -> dict:
    """Construct a structured GXRuntimeError dict.

    Args:
        error_type: Error class (e.g. "DecodeError", "LaneError", "SegmentError")
        message: Human-readable message
        segment_id: Optional segment id if segment-specific
        lane: Optional lane id if lane-specific
        recoverable: Whether error is recoverable

    Returns:
        GXRuntimeError dict
    """
    return {
        "type": error_type,
        "message": message,
        "segment_id": segment_id,
        "lane": lane,
        "recoverable": recoverable,
    }