"""
Basic MRS Reasoning Example

Demonstrates a 3-step reasoning chain: R₁ → R₂ → R₃
with drift computation and halting when C(Sₙ) ≥ τ

Run with: python -m mrs.examples.basic_reasoning
"""

import sys
from typing import Any
from mrs.core.recursion_node import RecursionNode
from mrs.core.state_tracker import StateTracker
from mrs.core.constraint_monitor import ConstraintMonitor
from mrs.core.topology_engine import TopologyEngine, TopologyMode
from mrs.utils.drift import compute_drift


def simple_reasoning_transform(input_text: Any) -> str:
    """
    Simple mock transformation simulating LLM reasoning.
    In practice, this would call an actual LLM.
    """
    text = str(input_text)

    if "decompose" in text.lower() or "step 1" in text.lower():
        return f"[Analysis] Breaking down: {text[:50]}... into sub-components"
    elif "analyze" in text.lower() or "step 2" in text.lower():
        return f"[Synthesis] Combining analysis results into coherent understanding"
    elif "synthesize" in text.lower() or "step 3" in text.lower():
        return f"[Conclusion] Final integrated answer based on analysis"
    else:
        return f"[Processed] Reasoning about: {text[:50]}..."


class MRS:
    """
    Minimal MRS implementation combining all components.
    """

    def __init__(
        self,
        model_fn=None,
        num_slots: int = 4,
        max_depth: int = 8,
        drift_threshold: float = 0.65,
    ):
        self.model_fn = model_fn or simple_reasoning_transform
        self.state = StateTracker(num_slots=num_slots)
        self.monitor = ConstraintMonitor(threshold=drift_threshold)
        self.topology = TopologyEngine(max_depth=max_depth, model_fn=self.model_fn)

    def run(self, prompt: str, verbose: bool = True) -> dict:
        """
        Execute MRS reasoning chain.

        Args:
            prompt: Initial query/prompt
            verbose: Whether to print step-by-step output

        Returns:
            Dictionary with history, slots, and final answer
        """
        self.state.reset()
        self.monitor.reset()
        self.topology.reset()

        history = []

        if verbose:
            print("\n" + "=" * 60)
            print("MRS v1.1 - Modular Reasoning Scaffold")
            print("=" * 60)
            print(f"\n[INPUT] {prompt}\n")
            print("-" * 60)

        current_node = self.topology.create_initial_node(prompt)

        while True:
            depth = self.topology.get_depth() + 1

            output = current_node.apply(self.model_fn)

            previous_output = self.state.get_last_output()
            drift = compute_drift(output, previous_output, method="hash")
            current_node.set_drift(drift)

            self.state.append(current_node)
            self.state.update_slot(depth, output)

            history.append(
                {
                    "depth": depth,
                    "input": str(current_node.x)[:100],
                    "output": str(output)[:100],
                    "drift": drift,
                }
            )

            total_drift = self.monitor.compute_drift(self.state)
            self.monitor.record_drift(total_drift)

            if verbose:
                print(f"[R{current_node.node_id}] Depth {depth}")
                print(f"    Input:  {str(current_node.x)[:60]}...")
                print(f"    Output: {str(output)[:60]}...")
                print(
                    f"    Δ = {drift:.4f}  |  C(S) = {total_drift:.4f}  |  τ = {self.monitor.threshold}"
                )

                if total_drift >= self.monitor.threshold:
                    print(f"    [HALT] Drift threshold exceeded: C(S) ≥ τ")
                elif drift < 0.05:
                    print(f"    [CONVERGED] Minimal drift detected")
                print()

            next_node = self.topology.step(current_node, self.state, self.monitor)

            if next_node is None:
                break

            current_node = next_node

        final_output = self.state.get_last_output()

        if verbose:
            print("-" * 60)
            print(f"\n[FINAL ANSWER]")
            print(f"  {final_output}")
            print(f"\n[STATS]")
            print(f"  Total steps: {len(history)}")
            print(f"  Final drift: C(S) = {self.monitor._drift_history[-1]:.4f}")
            print(f"  Mode: {self.topology.get_mode().value}")
            print("=" * 60 + "\n")

        return {
            "history": history,
            "slots": self.state._slots,
            "final": final_output,
            "total_drift": (
                self.monitor._drift_history[-1] if self.monitor._drift_history else 0
            ),
            "steps": len(history),
        }


def main():
    """CLI entrypoint for basic_reasoning example."""

    prompt = "Explain the concept of gravity: Step 1 decompose, Step 2 analyze, Step 3 synthesize"

    if len(sys.argv) > 1:
        prompt = " ".join(sys.argv[1:])

    mrs = MRS(
        model_fn=simple_reasoning_transform,
        num_slots=4,
        max_depth=8,
        drift_threshold=2.5,  # Higher threshold to allow multi-step chain
    )

    result = mrs.run(prompt, verbose=True)

    print("\n[DRIFT TRACE]")
    for step in result["history"]:
        bar_len = int(step["drift"] * 40)
        bar = "█" * bar_len + "░" * (40 - bar_len)
        print(f"  R{step['depth']}: [{bar}] Δ={step['drift']:.4f}")

    return result


if __name__ == "__main__":
    main()
