"""Tests for DeepSigmaCallbackHandler."""
import uuid

import pytest

from langchain_deepsigma import DeepSigmaCallbackHandler


def _run_id() -> uuid.UUID:
    return uuid.uuid4()


class TestDeepSigmaCallbackHandler:
    def test_init_defaults(self):
        handler = DeepSigmaCallbackHandler()
        assert handler.decision_type == "LangChainExecution"
        assert handler.emit_drift is True
        assert handler.verbose is False

    def test_chain_start_creates_episode(self):
        handler = DeepSigmaCallbackHandler(decision_type="TestChain")
        rid = _run_id()
        handler.on_chain_start(
            serialized={"name": "MyChain"},
            inputs={"input": "hello"},
            run_id=rid,
        )
        ep = handler.episode
        assert ep["decisionType"] == "MyChain"
        assert ep["outcome"]["code"] == "pending"
        assert "episodeId" in ep

    def test_chain_end_seals_episode_success(self):
        handler = DeepSigmaCallbackHandler()
        rid = _run_id()
        handler.on_chain_start(
            serialized={"name": "Chain"}, inputs={}, run_id=rid
        )
        handler.on_chain_end(outputs={"output": "ok"}, run_id=rid)

        ep = handler.episode
        assert ep["outcome"]["code"] == "success"
        assert "sealHash" in ep["seal"]
        assert ep["seal"]["version"] == 1

    def test_chain_end_runs_dlr_pipeline(self):
        handler = DeepSigmaCallbackHandler()
        rid = _run_id()
        handler.on_chain_start(
            serialized={"name": "Chain"}, inputs={}, run_id=rid
        )
        handler.on_chain_end(outputs={}, run_id=rid)

        # DLRBuilder should have entries
        assert len(handler.dlr.to_dict_list()) == 1

    def test_chain_error_outcome_fail_and_drift(self):
        handler = DeepSigmaCallbackHandler(emit_drift=True)
        rid = _run_id()
        handler.on_chain_start(
            serialized={"name": "Chain"}, inputs={}, run_id=rid
        )
        handler.on_chain_error(error=ValueError("boom"), run_id=rid)

        ep = handler.episode
        assert ep["outcome"]["code"] == "fail"
        assert len(handler.drift_events) == 1
        assert handler.drift_events[0]["driftType"] == "outcome"
        assert handler.drift_events[0]["severity"] == "red"

    def test_llm_start_adds_claim(self):
        handler = DeepSigmaCallbackHandler()
        rid = _run_id()
        handler.on_chain_start(serialized={}, inputs={}, run_id=rid)
        llm_rid = _run_id()
        handler.on_llm_start(
            serialized={}, prompts=["What is 2+2?"], run_id=llm_rid
        )
        claims = handler.episode["truth"]["claims"]
        assert len(claims) == 1
        assert "2+2" in claims[0]["statement"]
        assert claims[0]["statusLight"] == "GREEN"

    def test_llm_error_sets_red_claim(self):
        handler = DeepSigmaCallbackHandler(emit_drift=True)
        rid = _run_id()
        handler.on_chain_start(serialized={}, inputs={}, run_id=rid)
        llm_rid = _run_id()
        handler.on_llm_start(serialized={}, prompts=["hello"], run_id=llm_rid)
        handler.on_llm_error(error=RuntimeError("timeout"), run_id=llm_rid)

        claims = handler.episode["truth"]["claims"]
        assert claims[0]["statusLight"] == "RED"
        assert claims[0]["confidence"] == 0.0
        # drift emitted
        assert any(d["driftType"] == "verify" for d in handler.drift_events)

    def test_tool_start_adds_action(self):
        handler = DeepSigmaCallbackHandler()
        rid = _run_id()
        handler.on_chain_start(serialized={}, inputs={}, run_id=rid)
        tool_rid = _run_id()
        handler.on_tool_start(
            serialized={"name": "search"}, input_str="langchain", run_id=tool_rid
        )
        actions = handler.episode["actions"]
        assert len(actions) == 1
        assert actions[0]["type"] == "search"

    def test_tool_end_updates_action(self):
        handler = DeepSigmaCallbackHandler()
        rid = _run_id()
        handler.on_chain_start(serialized={}, inputs={}, run_id=rid)
        tool_rid = _run_id()
        handler.on_tool_start(
            serialized={"name": "calculator"}, input_str="2+2", run_id=tool_rid
        )
        handler.on_tool_end(output="4", run_id=tool_rid)

        action = handler.episode["actions"][0]
        assert action["result"] == "success"
        assert action["output"] == "4"

    def test_memory_graph_receives_episode(self):
        handler = DeepSigmaCallbackHandler()
        rid = _run_id()
        handler.on_chain_start(serialized={"name": "MG"}, inputs={}, run_id=rid)
        handler.on_chain_end(outputs={}, run_id=rid)

        assert handler.mg.node_count >= 1
