"""Tests for Super Big Brain: atomic decomposition, graph, recomposition."""

import pytest

from fusionagi.schemas.atomic import (
    AtomicSemanticUnit,
    AtomicUnitType,
    DecompositionResult,
    SemanticRelation,
    RelationType,
)
from fusionagi.reasoning.decomposition import decompose_recursive
from fusionagi.memory.semantic_graph import SemanticGraphMemory
from fusionagi.memory.sharding import shard_context, Shard
from fusionagi.reasoning.context_loader import load_context_for_reasoning, build_compact_prompt
from fusionagi.memory.scratchpad import LatentScratchpad, ThoughtState
from fusionagi.reasoning.tot import ThoughtNode, expand_node, prune_subtree, merge_subtrees
from fusionagi.reasoning.multi_path import generate_and_score_parallel
from fusionagi.reasoning.recomposition import recompose, RecomposedResponse
from fusionagi.reasoning.meta_reasoning import challenge_assumptions, detect_contradictions, revisit_node
from fusionagi.core.super_big_brain import (
    run_super_big_brain,
    SuperBigBrainConfig,
    SuperBigBrainReasoningProvider,
)
from fusionagi.schemas.head import HeadId


class TestAtomicSchema:
    """Test atomic semantic unit schemas."""

    def test_atomic_unit_creation(self):
        u = AtomicSemanticUnit(
            unit_id="asu_1",
            content="Test fact",
            type=AtomicUnitType.FACT,
            confidence=0.9,
        )
        assert u.unit_id == "asu_1"
        assert u.content == "Test fact"
        assert u.type == AtomicUnitType.FACT
        assert u.confidence == 0.9

    def test_decomposition_result(self):
        u = AtomicSemanticUnit(unit_id="asu_1", content="Fact", type=AtomicUnitType.FACT)
        r = SemanticRelation(from_id="root", to_id="asu_1", relation_type=RelationType.LOGICAL)
        result = DecompositionResult(units=[u], relations=[r], depth=0)
        assert len(result.units) == 1
        assert len(result.relations) == 1
        assert result.depth == 0


class TestDecomposition:
    """Test recursive decomposition."""

    def test_decompose_simple(self):
        result = decompose_recursive("What are the security risks? Must support 1M users.")
        assert len(result.units) >= 1
        assert result.depth >= 0

    def test_decompose_empty(self):
        result = decompose_recursive("")
        assert len(result.units) == 0

    def test_decompose_max_depth(self):
        result = decompose_recursive("Question one? Question two? Question three?", max_depth=1)
        assert result.depth <= 1


class TestSemanticGraph:
    """Test semantic graph memory."""

    def test_add_and_query(self):
        g = SemanticGraphMemory()
        u = AtomicSemanticUnit(unit_id="asu_1", content="Fact", type=AtomicUnitType.FACT)
        g.add_unit(u)
        assert g.get_unit("asu_1") == u
        assert len(g.query_units()) >= 1

    def test_ingest_decomposition(self):
        g = SemanticGraphMemory()
        r = decompose_recursive("What is X? Constraint: must be fast.")
        g.ingest_decomposition(r.units, r.relations)
        assert len(g.query_units()) >= 1


class TestSharding:
    """Test context sharding."""

    def test_shard_context(self):
        r = decompose_recursive("Security risk? Cost constraint?")
        shards = shard_context(r.units, max_cluster_size=5)
        assert isinstance(shards, list)
        assert all(isinstance(s, Shard) for s in shards)


class TestContextLoader:
    """Test retrieve-by-reference."""

    def test_load_context(self):
        r = decompose_recursive("Test prompt")
        ctx = load_context_for_reasoning(r.units)
        assert "unit_refs" in ctx
        assert "unit_summaries" in ctx

    def test_build_compact_prompt(self):
        r = decompose_recursive("Short prompt")
        prompt = build_compact_prompt(r.units, max_chars=1000)
        assert isinstance(prompt, str)


class TestScratchpad:
    """Test latent scratchpad."""

    def test_append_and_get(self):
        s = LatentScratchpad()
        s.append_hypothesis("H1")
        s.append_discarded("D1")
        state = s.get_intermediate()
        assert len(state.hypotheses) == 1
        assert len(state.discarded_paths) == 1

    def test_clear(self):
        s = LatentScratchpad()
        s.append_hypothesis("H1")
        s.clear()
        state = s.get_intermediate()
        assert len(state.hypotheses) == 0


class TestMetaReasoning:
    """Test meta-reasoning hooks."""

    def test_challenge_assumptions(self):
        u = AtomicSemanticUnit(
            unit_id="asu_1",
            content="Assume X is true",
            type=AtomicUnitType.ASSUMPTION,
        )
        flagged = challenge_assumptions([u], "Conclusion based on X")
        assert len(flagged) >= 0

    def test_detect_contradictions(self):
        u1 = AtomicSemanticUnit(unit_id="a", content="X is true", type=AtomicUnitType.FACT)
        u2 = AtomicSemanticUnit(unit_id="b", content="X is not true", type=AtomicUnitType.FACT)
        pairs = detect_contradictions([u1, u2])
        assert isinstance(pairs, list)


class TestSuperBigBrain:
    """Test Super Big Brain orchestrator."""

    def test_run_super_big_brain(self):
        g = SemanticGraphMemory()
        r = run_super_big_brain("What are the risks?", g)
        assert isinstance(r, RecomposedResponse)
        assert r.summary
        assert 0 <= r.confidence <= 1

    def test_super_big_brain_reasoning_provider(self):
        p = SuperBigBrainReasoningProvider()
        ho = p.produce_head_output(HeadId.LOGIC, "Analyze architecture")
        assert ho.head_id == HeadId.LOGIC
        assert ho.summary
        assert len(ho.claims) >= 0