#!/usr/bin/env python3
"""
Echo Bridge: Connect cognitive tasks to BloomOS/Helix64.

This module bridges Echo-style cognitive tasks (Tower of Hanoi, Blocks World, etc.)
to the BloomOS runtime, enabling Helix64 to serve as the "CPU" for autonomous
cognitive reasoning.

Architecture:
    Cog Task JSON → EchoBridge → BloomOS Task → Helix64 VM → Response → State Update

Usage:
    from helix_cdc.echo_bridge import EchoBridge, CogTask

    # Load a cognitive task
    bridge = EchoBridge(bloomos_url="http://localhost:7070")
    task = bridge.load_task("path/to/SE_HANOI.json")

    # Run a single reasoning step
    result = bridge.step(task, current_state)

    # Run full task loop
    final_state, trace = bridge.run_task(task, initial_state, max_steps=20)
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from helix_cdc.bloomos.bloomos_client import BloomOSClient
from helix_cdc.bloomos.bloomos_agent import BloomOSAgent, AgentStep
from helix_cdc.bloomos.sym import SymCapsule, lower_task_to_sym
from helix_cdc.bloomos.pheromone import PheromoneBoard, make_entry_from_step
from helix_cdc.metrics import measure_text, SeMetrics
from helix_cdc.seed_api import encode_text_to_seed, regrow_seed, compute_seed_entropy, get_dynamic_lattice_D


@dataclass
class CogTaskState:
    """Represents the current state of a cognitive task."""
    state_repr: str  # String representation of current state
    step_count: int = 0
    history: List[str] = field(default_factory=list)
    is_goal: bool = False
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class CogTask:
    """
    A cognitive task definition.

    Loaded from JSON files like SE_HANOI.json, SE_BLOCKS.json, etc.
    """
    task_id: str
    task_type: str  # "hanoi", "blocks", "reasoning", etc.
    description: str
    initial_state: str
    goal_state: str
    rules: List[str] = field(default_factory=list)
    hints: List[str] = field(default_factory=list)
    difficulty: str = "medium"
    max_steps: int = 50
    metadata: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "CogTask":
        """Create a CogTask from JSON data."""
        return cls(
            task_id=data.get("task_id", data.get("id", "unknown")),
            task_type=data.get("task_type", data.get("type", "reason")),
            description=data.get("description", ""),
            initial_state=data.get("initial_state", data.get("start", "")),
            goal_state=data.get("goal_state", data.get("goal", "")),
            rules=data.get("rules", []),
            hints=data.get("hints", []),
            difficulty=data.get("difficulty", "medium"),
            max_steps=data.get("max_steps", 50),
            metadata=data.get("metadata", {}),
        )

    @classmethod
    def from_file(cls, path: Union[str, Path]) -> "CogTask":
        """Load a CogTask from a JSON file."""
        path = Path(path)
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        task = cls.from_json(data)
        task.metadata["source_file"] = str(path)
        return task


@dataclass
class StepResult:
    """Result of a single cognitive step."""
    action: str
    reasoning: str
    new_state: str
    is_goal: bool
    logits_mean: Optional[float]
    logits_std: Optional[float]
    status: str
    raw_response: Dict[str, Any] = field(default_factory=dict)
    elapsed_ms: float = 0.0
    # Symbolic Entropy metrics (Se = H * C * D)
    se_H: Optional[float] = None  # Shannon entropy (bits/byte)
    se_C: Optional[float] = None  # Contextual coherence (0-1)
    se_D: Optional[float] = None  # Dimensional depth (1-72)
    se_Se: Optional[float] = None  # Symbolic entropy = H * C * D
    # SymCapsule for the step
    capsule: Optional[SymCapsule] = None
    # CDR seed info (optional)
    seed_len: Optional[int] = None
    seed_entropy: Optional[float] = None
    regrow_depth: Optional[int] = None
    lattice_D: Optional[int] = None


class EchoBridge:
    """
    Bridge between Echo cognitive tasks and BloomOS/Helix64.

    This class orchestrates multi-step reasoning tasks through Helix64,
    turning the VM into a cognitive "CPU" that can solve problems like
    Tower of Hanoi, Blocks World, etc.
    """

    def __init__(
        self,
        bloomos_url: str = "http://localhost:7070",
        timeout: float = 60.0,
        helix_program: Optional[str] = None,
        log_path: Optional[str] = None,
        trace_path: Optional[str] = None,
        pheromone_path: Optional[str] = None,
        regrow_depth: Optional[int] = None,
    ):
        """
        Initialize the Echo bridge.

        Args:
            bloomos_url: BloomOS runtime URL
            timeout: HTTP timeout for requests
            helix_program: Optional Helix program to use (e.g., "BALANCED64")
            log_path: Optional path for JSONL logging (BloomOS client)
            trace_path: Optional path for cognitive trace JSONL (SymCapsule + Se per step)
            pheromone_path: Optional path for pheromone board JSONL (adaptive learning signals)
            regrow_depth: Optional CDR seed regrow depth (0=disabled, 64/256/1024=enabled)
        """
        self.client = BloomOSClient(
            base_url=bloomos_url,
            timeout=timeout,
            log_path=log_path,
        )
        self.agent = BloomOSAgent(client=self.client)
        self.helix_program = helix_program
        self.trace_path = Path(trace_path) if trace_path else None
        self.pheromone_board = PheromoneBoard(pheromone_path) if pheromone_path else None
        self.regrow_depth = regrow_depth or 0
        self.task_type_map = {
            "hanoi": "reason",
            "blocks": "reason",
            "reasoning": "reason",
            "planning": "reason",
            "math": "math",
            "code": "code",
            "json": "json",
        }

    def load_task(self, path: Union[str, Path]) -> CogTask:
        """Load a cognitive task from a JSON file."""
        return CogTask.from_file(path)

    def _write_cognitive_trace(
        self,
        task: CogTask,
        step_num: int,
        step_result: StepResult,
        current_state: CogTaskState,
    ) -> None:
        """
        Write a cognitive trace entry with SymCapsule + Se metrics.

        Args:
            task: The cognitive task definition
            step_num: Current step number (0-indexed)
            step_result: Result from this step
            current_state: State before this step
        """
        if not self.trace_path:
            return

        # Build trace entry
        entry = {
            "timestamp": time.time(),
            "timestamp_iso": time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime()),
            "task_id": task.task_id,
            "task_type": task.task_type,
            "step_num": step_num,
            "state_before": current_state.state_repr,
            "state_after": step_result.new_state,
            "action": step_result.action,
            "reasoning": step_result.reasoning,
            "is_goal": step_result.is_goal,
            "status": step_result.status,
            "elapsed_ms": step_result.elapsed_ms,
            "logits_mean": step_result.logits_mean,
            "logits_std": step_result.logits_std,
            # Se metrics
            "se_metrics": {
                "H": step_result.se_H,
                "C": step_result.se_C,
                "D": step_result.se_D,
                "Se": step_result.se_Se,
            },
            # SymCapsule (serialized)
            "capsule": step_result.capsule.to_dict() if step_result.capsule else None,
        }

        try:
            with self.trace_path.open("a", encoding="utf-8") as f:
                f.write(json.dumps(entry) + "\n")
        except OSError:
            pass  # Fail silently

    def _build_prompt(
        self,
        task: CogTask,
        current_state: CogTaskState,
    ) -> str:
        """
        Build a prompt for a single reasoning step.

        Args:
            task: The cognitive task definition
            current_state: Current state of the task

        Returns:
            Prompt string for BloomOS
        """
        lines = [
            f"Task: {task.description}",
            "",
            f"Current State: {current_state.state_repr}",
            f"Goal State: {task.goal_state}",
            "",
        ]

        if task.rules:
            lines.append("Rules:")
            for r in task.rules:
                lines.append(f"  - {r}")
            lines.append("")

        if current_state.history:
            lines.append(f"Previous actions ({len(current_state.history)}):")
            for i, action in enumerate(current_state.history[-5:], 1):
                lines.append(f"  {i}. {action}")
            lines.append("")

        lines.extend([
            "Analyze the current state and determine the next action.",
            "Respond with:",
            "1. Your reasoning (1-2 sentences)",
            "2. The action to take",
            "3. The resulting state after the action",
        ])

        return "\n".join(lines)

    def _parse_response(
        self,
        response: Dict[str, Any],
        current_state: CogTaskState,
        goal_state: str,
    ) -> StepResult:
        """
        Parse BloomOS response into a StepResult.

        Args:
            response: Raw BloomOS response
            current_state: Current state before this step
            goal_state: Target goal state

        Returns:
            StepResult with parsed action, reasoning, and new state
        """
        # Extract output
        output = response.get("output", {})
        status = response.get("status", "error")
        logits_mean = response.get("logits_mean")
        logits_std = response.get("logits_std")

        # Default values
        action = "no_action"
        reasoning = "Unable to determine action"
        new_state = current_state.state_repr

        # Parse the response if available
        # (In a real implementation, this would parse LLM output)
        if status == "ok" and output:
            # For now, return placeholder - actual parsing depends on output format
            action = output.get("action", "step")
            reasoning = output.get("reasoning", "Computed next step")
            new_state = output.get("new_state", current_state.state_repr)

        # Check if we reached the goal
        is_goal = new_state.strip() == goal_state.strip()

        return StepResult(
            action=action,
            reasoning=reasoning,
            new_state=new_state,
            is_goal=is_goal,
            logits_mean=logits_mean,
            logits_std=logits_std,
            status=status,
            raw_response=response,
        )

    def step(
        self,
        task: CogTask,
        current_state: CogTaskState,
        use_agent: bool = False,
    ) -> StepResult:
        """
        Execute a single reasoning step.

        Args:
            task: The cognitive task definition
            current_state: Current state of the task
            use_agent: If True, use BloomOSAgent.classify_goal() for auto-routing

        Returns:
            StepResult with action, reasoning, and new state
        """
        # Build prompt
        prompt = self._build_prompt(task, current_state)

        # Execute via BloomOS
        t0 = time.perf_counter()

        if use_agent:
            # Use BloomOSAgent for intelligent task classification
            agent_step: AgentStep = self.agent.classify_goal(prompt)
            bloomos_task_type = agent_step.task_type
            gate_policy = agent_step.gate_policy
            anchors = agent_step.anchors
        else:
            # Map task type to BloomOS task type using static map
            bloomos_task_type = self.task_type_map.get(task.task_type, "reason")
            gate_policy = None
            anchors = None

        extra = {}
        if self.helix_program:
            extra["helix_program"] = self.helix_program

        result = self.client.run_task(
            task_type=bloomos_task_type,
            description=prompt,
            anchors=anchors,
            gate_policy=gate_policy,
            extra=extra if extra else None,
        )

        elapsed_ms = (time.perf_counter() - t0) * 1000.0

        # Parse response
        step_result = self._parse_response(
            result.raw if hasattr(result, "raw") else {},
            current_state,
            task.goal_state,
        )
        step_result.elapsed_ms = elapsed_ms

        # Compute Se metrics from output
        output_text = self._extract_output_text(result)
        if output_text:
            try:
                se_metrics = measure_text(output_text, auto_detect_d=True)
                step_result.se_H = se_metrics.H
                step_result.se_C = se_metrics.C
                step_result.se_D = se_metrics.D
                step_result.se_Se = se_metrics.Se
            except Exception:
                pass  # Se computation failed, leave as None

        # Create SymCapsule for this step
        helix_program = None
        try:
            helix_program = result.raw.get("helix_program", "ABCDEFGH")
            step_result.capsule = lower_task_to_sym(
                task_type=bloomos_task_type,
                description=prompt,
                program=helix_program,
                result=result.raw,
            )
        except Exception:
            pass  # Capsule creation failed, leave as None

        # Record pheromone entry for adaptive learning
        if self.pheromone_board is not None:
            try:
                # Get lattice D for enriched pheromone signal
                lattice_d = get_dynamic_lattice_D()
                step_result.lattice_D = lattice_d

                pheromone_entry = make_entry_from_step(
                    task_type=bloomos_task_type,
                    gate=gate_policy,
                    helix_program=helix_program,
                    status=step_result.status,
                    logits_mean=step_result.logits_mean,
                    logits_std=step_result.logits_std,
                    se_H=step_result.se_H,
                    se_C=step_result.se_C,
                    se_D=step_result.se_D,
                    se_Se=step_result.se_Se,
                    meta={
                        "task_id": task.task_id if hasattr(task, "task_id") else None,
                        "elapsed_ms": step_result.elapsed_ms,
                    },
                    # Phase 5 enriched fields
                    seed_len=step_result.seed_len,
                    seed_entropy=step_result.seed_entropy,
                    regrow_depth=step_result.regrow_depth,
                    lattice_D=lattice_d,
                )
                self.pheromone_board.record(pheromone_entry)
            except Exception:
                pass  # Pheromone recording failed, non-fatal

        # Optional: CDR seed encoding & regrow
        if self.regrow_depth and output_text:
            try:
                seed_info = encode_text_to_seed(output_text)
                step_result.seed_len = len(seed_info.seed_bytes)
                step_result.seed_entropy = seed_info.seed_entropy
                step_result.regrow_depth = self.regrow_depth
                # Regrow the seed (result ignored for now, just exercising the pipeline)
                _ = regrow_seed(seed_info.seed_bytes, depth=self.regrow_depth)
            except Exception:
                # Non-fatal; leave seed fields as None
                pass

        return step_result

    def _extract_output_text(self, result) -> Optional[str]:
        """Extract text from TaskResult output for Se measurement."""
        output = result.output if hasattr(result, "output") else None
        if output is None:
            return None
        if isinstance(output, str):
            return output
        if isinstance(output, dict):
            for key in ("text", "content", "response", "output", "result", "reasoning"):
                if key in output and isinstance(output[key], str):
                    return output[key]
            return json.dumps(output)
        return str(output)

    def run_task(
        self,
        task: CogTask,
        initial_state: Optional[CogTaskState] = None,
        max_steps: Optional[int] = None,
        verbose: bool = True,
        use_agent: bool = False,
    ) -> tuple[CogTaskState, List[StepResult]]:
        """
        Run a full cognitive task loop until goal or max steps.

        Args:
            task: The cognitive task definition
            initial_state: Starting state (defaults to task.initial_state)
            max_steps: Maximum steps (defaults to task.max_steps)
            verbose: Print progress
            use_agent: If True, use BloomOSAgent for intelligent task routing

        Returns:
            (final_state, trace) tuple
        """
        max_steps = max_steps or task.max_steps

        # Initialize state
        if initial_state is None:
            current_state = CogTaskState(
                state_repr=task.initial_state,
                step_count=0,
                history=[],
                is_goal=False,
            )
        else:
            current_state = initial_state

        trace: List[StepResult] = []

        if verbose:
            print(f"=== Running Cognitive Task: {task.task_id} ===")
            print(f"Description: {task.description}")
            print(f"Initial: {current_state.state_repr}")
            print(f"Goal: {task.goal_state}")
            print(f"Max steps: {max_steps}")
            print()

        for step_num in range(max_steps):
            if verbose:
                print(f"Step {step_num + 1}: ", end="", flush=True)

            # Execute step
            state_before = current_state  # Save for trace logging
            result = self.step(task, current_state, use_agent=use_agent)
            trace.append(result)

            # Write cognitive trace entry (SymCapsule + Se per step)
            self._write_cognitive_trace(task, step_num, result, state_before)

            if verbose:
                se_str = f", Se={result.se_Se:.2f}" if result.se_Se else ""
                print(f"{result.action} (logits={result.logits_mean:.4f}{se_str})"
                      if result.logits_mean else f"{result.action}{se_str}")

            # Update state
            current_state = CogTaskState(
                state_repr=result.new_state,
                step_count=step_num + 1,
                history=current_state.history + [result.action],
                is_goal=result.is_goal,
            )

            # Check goal
            if result.is_goal:
                if verbose:
                    print(f"\nGoal reached in {step_num + 1} steps!")
                break

            # Check for stuck state
            if result.status != "ok":
                if verbose:
                    print(f"\nError: {result.status}")
                break

        return current_state, trace

    def create_task(
        self,
        task_type: str,
        description: str,
        initial_state: str,
        goal_state: str,
        rules: Optional[List[str]] = None,
        task_id: Optional[str] = None,
    ) -> CogTask:
        """
        Create a cognitive task programmatically.

        Args:
            task_type: Type of task (hanoi, blocks, reasoning, etc.)
            description: Human-readable description
            initial_state: Starting state representation
            goal_state: Target state representation
            rules: Optional list of rules/constraints
            task_id: Optional task identifier

        Returns:
            CogTask instance
        """
        return CogTask(
            task_id=task_id or f"{task_type}_{int(time.time())}",
            task_type=task_type,
            description=description,
            initial_state=initial_state,
            goal_state=goal_state,
            rules=rules or [],
        )


# Example cognitive tasks for testing
EXAMPLE_TASKS = {
    "hanoi_3": {
        "task_id": "hanoi_3",
        "task_type": "hanoi",
        "description": "Solve Tower of Hanoi with 3 disks",
        "initial_state": "A:[3,2,1] B:[] C:[]",
        "goal_state": "A:[] B:[] C:[3,2,1]",
        "rules": [
            "Move one disk at a time",
            "Cannot place larger disk on smaller disk",
            "Only move top disk from a peg",
        ],
        "max_steps": 15,
    },
    "blocks_simple": {
        "task_id": "blocks_simple",
        "task_type": "blocks",
        "description": "Stack blocks in correct order",
        "initial_state": "Table:[A,B,C] (all on table)",
        "goal_state": "Stack:[C on B on A on table]",
        "rules": [
            "Move one block at a time",
            "Can only move a block with nothing on top",
            "Blocks can be on table or on another block",
        ],
        "max_steps": 10,
    },
}


def get_example_task(name: str) -> CogTask:
    """Get a built-in example task by name."""
    if name not in EXAMPLE_TASKS:
        available = list(EXAMPLE_TASKS.keys())
        raise ValueError(f"Unknown example task '{name}'. Available: {available}")
    return CogTask.from_json(EXAMPLE_TASKS[name])


# CLI entry point
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Run cognitive tasks through Helix64/BloomOS"
    )
    parser.add_argument(
        "--task",
        help="Path to task JSON file or example name (hanoi_3, blocks_simple)",
    )
    parser.add_argument(
        "--bloomos-url",
        default="http://localhost:7070",
        help="BloomOS runtime URL",
    )
    parser.add_argument(
        "--max-steps",
        type=int,
        help="Maximum steps (overrides task default)",
    )
    parser.add_argument(
        "--helix-program",
        help="Helix program name or sequence",
    )
    parser.add_argument(
        "--list-examples",
        action="store_true",
        help="List available example tasks",
    )
    parser.add_argument(
        "--use-agent",
        action="store_true",
        help="Use BloomOSAgent for intelligent task routing",
    )
    parser.add_argument(
        "--trace-path",
        help="Path for cognitive trace JSONL (SymCapsule + Se per step)",
    )
    parser.add_argument(
        "--pheromone-path",
        help="Path for pheromone board JSONL (adaptive learning signals)",
    )
    parser.add_argument(
        "--regrow-depth",
        type=int,
        default=0,
        help="If >0, encode each step's output as CDR seed and regrow to this depth.",
    )

    args = parser.parse_args()

    if args.list_examples:
        print("Available example tasks:")
        for name, task_data in EXAMPLE_TASKS.items():
            print(f"  {name}: {task_data['description']}")
        exit(0)

    if not args.task:
        parser.print_help()
        exit(1)

    # Initialize bridge
    bridge = EchoBridge(
        bloomos_url=args.bloomos_url,
        helix_program=args.helix_program,
        trace_path=args.trace_path,
        pheromone_path=args.pheromone_path,
        regrow_depth=args.regrow_depth,
    )

    # Load task
    if args.task in EXAMPLE_TASKS:
        task = get_example_task(args.task)
    else:
        task = bridge.load_task(args.task)

    # Run task
    final_state, trace = bridge.run_task(
        task,
        max_steps=args.max_steps,
        verbose=True,
        use_agent=args.use_agent,
    )

    # Summary
    print()
    print("=== Summary ===")
    print(f"Final state: {final_state.state_repr}")
    print(f"Goal reached: {final_state.is_goal}")
    print(f"Steps taken: {final_state.step_count}")
    print(f"Total trace entries: {len(trace)}")
