from .emma_mdp import EmmaState, EmmaAction, EmmaContext, TRANSITIONS, reward
from .emma_resonance import MorphicMemory
from .emma_flynn import FlynnTracker


class EmmaMCoTAgent:
    """
    Markov Chain of Thought agent loop (NAACL 2025).
    Kombinerer MDP, Morphic Resonance og Flynn-tracker.
    """

    def __init__(self, llm, memory: MorphicMemory, mdp_transitions=None):
        self.llm = llm
        self.memory = memory
        self.transitions = mdp_transitions or TRANSITIONS
        self.state = EmmaState.IDLE
        self.flynn_tracker = FlynnTracker()

    def run(self, userinput: str, history: list) -> str:
        ctx = EmmaContext(userinput=userinput, history=history)
        ctx.complexityscore = self.flynn_tracker.measure(userinput)
        self.state = EmmaState.OBSERVING

        while self.state != EmmaState.IDLE:
            embedding = self.llm.embed(str(ctx))
            patterns = self.memory.recall(embedding)
            thought = self.llm.think(ctx, patterns, self.state)
            ctx.thoughtchain.append(thought)
            ctx.stepcount += 1

            if len(ctx.thoughtchain) > 4:
                ctx = self._compress_chain(ctx)

            action = self._policy(ctx, thought, patterns)
            outcome = self._execute(action, ctx)
            r = reward(self.state, action, outcome)
            ctx.rewardacc += r

            if r > 0.3:
                self.memory.store(embedding, action, r)

            next_state = self.transitions.get((self.state, action))
            if next_state is None:
                break
            self.state = next_state

        self.flynn_tracker.record(
            complexity=ctx.complexityscore,
            steps=ctx.stepcount,
            reward=ctx.rewardacc,
        )
        return ctx.thoughtchain[-1] if ctx.thoughtchain else ""

    def _compress_chain(self, ctx: EmmaContext) -> EmmaContext:
        summary = self.llm.compress(ctx.thoughtchain[:-1])
        ctx.thoughtchain = [summary, ctx.thoughtchain[-1]]
        return ctx

    def _policy(self, ctx, thought, patterns) -> EmmaAction:
        return self.llm.choose_action(ctx, thought, patterns, list(EmmaAction))

    def _execute(self, action: EmmaAction, ctx: EmmaContext) -> dict:
        try:
            result = self.llm.execute_action(action, ctx)
            return {"task_completed": True, "steps_used": ctx.stepcount, **result}
        except Exception as e:
            return {"tool_error": True, "error": str(e)}