# Copyright (c) 2023-2026, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

"""Integration tests for lionagi CLI orchestration paths — catches regressions unit tests miss."""

from __future__ import annotations

import asyncio
import inspect

import pytest

# ── Test 1: cancelled_exc_classes safe to call outside loop ──────────────────


def test_cancelled_exc_safe_outside_loop():
    """cancelled_exc_classes() must not raise when called from a sync context (no running event loop)."""
    from lionagi.ln.concurrency.errors import cancelled_exc_classes

    result = cancelled_exc_classes()

    assert isinstance(result, tuple)
    assert len(result) >= 1
    # asyncio.CancelledError must always be in the tuple (the safe baseline)
    assert asyncio.CancelledError in result


def test_is_cancelled_works_with_cancelled_error():
    """is_cancelled() must recognise asyncio.CancelledError outside a loop."""
    from lionagi.ln.concurrency.errors import is_cancelled

    exc = asyncio.CancelledError()
    assert is_cancelled(exc) is True


def test_is_cancelled_false_for_non_cancel():
    """is_cancelled() must return False for non-cancellation exceptions."""
    from lionagi.ln.concurrency.errors import is_cancelled

    assert is_cancelled(ValueError("not a cancel")) is False
    assert is_cancelled(RuntimeError()) is False


# ── Test 2: cache_cancelled_exc_class populates the cache ────────────────────


async def test_cancelled_exc_cache_populated_after_explicit_cache():
    """cache_cancelled_exc_class() inside an event loop populates the module cache."""
    from lionagi.ln.concurrency import errors as _err_mod
    from lionagi.ln.concurrency.errors import cache_cancelled_exc_class, cancelled_exc_classes

    # Reset module-level cache to simulate first call
    original = _err_mod._CANCELLED_EXC_CLASS
    _err_mod._CANCELLED_EXC_CLASS = None
    try:
        # Must be called from inside an event loop
        cache_cancelled_exc_class()
        result = cancelled_exc_classes()
        assert isinstance(result, tuple)
        assert len(result) >= 1
        assert asyncio.CancelledError in result
    finally:
        # Restore original state
        _err_mod._CANCELLED_EXC_CLASS = original


# ── Test 3: aggregation_sources in metadata, NOT in parameters ───────────────


def test_aggregation_params_in_metadata_not_parameters():
    """build_fanout_graph must place aggregation_sources/aggregation_count in metadata, not parameters."""
    from lionagi.casts.emission import TaskAssignment
    from lionagi.orchestration.patterns import build_fanout_graph
    from lionagi.session.branch import Branch
    from lionagi.session.session import Session

    session = Session()

    # Create two role branches and register them
    worker_a = Branch(name="worker_a")
    worker_b = Branch(name="worker_b")
    synth = Branch(name="synth")
    session.include_branches(worker_a)
    session.include_branches(worker_b)
    session.include_branches(synth)

    roles = {"worker_a": worker_a, "worker_b": worker_b, "synth": synth}

    assignments = [
        TaskAssignment(task="Do task A", assignee="worker_a"),
        TaskAssignment(task="Do task B", assignee="worker_b"),
    ]

    graph, worker_ids = build_fanout_graph(
        session,
        assignments,
        roles,
        synthesis_role="synth",
    )

    # Find the synthesis node (non-worker)
    from lionagi.operations.node import Operation

    synth_node = None
    for node in graph.internal_nodes.values():
        if isinstance(node, Operation) and node.metadata.get("aggregation"):
            synth_node = node
            break

    assert synth_node is not None, "No synthesis node found in graph"

    # aggregation keys must be in metadata
    assert "aggregation_sources" in synth_node.metadata, (
        "aggregation_sources must be in Operation.metadata, not parameters"
    )
    assert "aggregation_count" in synth_node.metadata, (
        "aggregation_count must be in Operation.metadata, not parameters"
    )
    assert len(synth_node.metadata["aggregation_sources"]) == 2

    # parameters must only contain instruction — NOT aggregation kwargs
    params = synth_node.parameters
    if isinstance(params, dict):
        assert "aggregation_sources" not in params, (
            "aggregation_sources must NOT be in parameters (causes TypeError in operate())"
        )
        assert "aggregation_count" not in params, (
            "aggregation_count must NOT be in parameters (causes TypeError in operate())"
        )
        assert "instruction" in params


# ── Test 4: _on_bus_spawn is async ───────────────────────────────────────────


def test_on_bus_spawn_is_async():
    """ReactiveExecutor._on_bus_spawn must be async so session.observe() emit gathers it as a coro."""
    from lionagi.operations.flow import ReactiveExecutor

    assert inspect.iscoroutinefunction(ReactiveExecutor._on_bus_spawn), (
        "ReactiveExecutor._on_bus_spawn must be a coroutine function (async def)"
    )


# ── Test 5: Session has hooks property returning HookBus ─────────────────────


def test_session_has_hooks_property():
    """Session.hooks must lazily return a HookBus bound to session.observer."""
    from lionagi.hooks import HookBus
    from lionagi.session.session import Session

    session = Session()

    # hooks is a lazy property — accessing it creates the bus
    bus = session.hooks

    assert isinstance(bus, HookBus), f"session.hooks must return HookBus, got {type(bus)}"
    assert session._hooks is not None, "session._hooks must be populated after first access"
    assert session._hooks is bus, "session._hooks and session.hooks must be the same object"

    # The bus must be bound to the session's observer
    assert bus._observer is session.observer, (
        "HookBus._observer must be session.observer (ADR-0076)"
    )


def test_session_hooks_identity_stable():
    """Accessing session.hooks multiple times returns the same bus."""
    from lionagi.hooks import HookBus
    from lionagi.session.session import Session

    session = Session()
    bus1 = session.hooks
    bus2 = session.hooks

    assert bus1 is bus2, "session.hooks must return the same HookBus instance on repeated access"


# ── Test 6: Branch gets hooks from session via include_branches ───────────────


def test_branch_gets_hooks_from_session():
    """include_branches must propagate session._hooks to each added branch."""
    from lionagi.hooks import HookBus
    from lionagi.session.branch import Branch
    from lionagi.session.session import Session

    session = Session()
    # Initialise the hook bus on the session BEFORE adding the branch
    _ = session.hooks  # forces _hooks creation

    new_branch = Branch(name="test_branch")
    session.include_branches(new_branch)

    assert new_branch._hooks is not None, "Branch._hooks must be set after include_branches"
    assert new_branch._hooks is session._hooks, (
        "branch._hooks must be the same object as session._hooks"
    )
    assert isinstance(new_branch._hooks, HookBus)


def test_branch_gets_hooks_when_added_after_bus_init():
    """Branches added after hooks init must receive the already-created bus."""
    from lionagi.session.branch import Branch
    from lionagi.session.session import Session

    session = Session()
    bus = session.hooks  # create bus first

    b1 = Branch(name="b1")
    b2 = Branch(name="b2")
    session.include_branches(b1)
    session.include_branches(b2)

    assert b1._hooks is bus
    assert b2._hooks is bus


# ── Test 7: both execute() and execute_stream() use the public observer property ──


def test_reactive_executor_uses_public_observer_property():
    """ReactiveExecutor.execute() must use session.observer (public property), not _observer.

    The private _observer PrivateAttr is None until first accessed; relying on
    it in execute() would silently drop reactive spawns when the observer hasn't
    been pre-initialised. Both execute() and execute_stream() must use the
    lazy-init public property instead.
    """
    source_execute = inspect.getsource(
        __import__(
            "lionagi.operations.flow", fromlist=["ReactiveExecutor"]
        ).ReactiveExecutor.execute
    )
    source_stream = inspect.getsource(
        __import__(
            "lionagi.operations.flow", fromlist=["ReactiveExecutor"]
        ).ReactiveExecutor.execute_stream
    )

    for method_name, source in (("execute", source_execute), ("execute_stream", source_stream)):
        # Must NOT access the private _observer attribute via getattr
        private_lines = [
            (i + 1, line)
            for i, line in enumerate(source.splitlines())
            if '"_observer"' in line or "'_observer'" in line
        ]
        assert not private_lines, (
            f"ReactiveExecutor.{method_name}() must not access session._observer "
            f"(private PrivateAttr). Found at lines: {private_lines}"
        )

        # Must use the public property (self.session.observer)
        public_lines = [
            line
            for line in source.splitlines()
            if "session.observer" in line and not line.strip().startswith("#")
        ]
        assert public_lines, (
            f"ReactiveExecutor.{method_name}() must access session.observer "
            f"(public lazy-init property) to ensure the observer is always initialised"
        )


# ── Test 8: _wait_for_dependencies reads aggregation_sources from metadata ────


def test_flow_aggregation_wait_reads_metadata():
    """DependencyAwareExecutor._wait_for_dependencies reads aggregation_sources
    from operation.metadata (not operation.parameters).

    If the executor read from parameters, it would find nothing (since
    build_fanout_graph stores them in metadata) and skip the aggregation wait,
    causing the synthesis to fire before workers complete.
    """
    from lionagi.operations.flow import DependencyAwareExecutor

    # Inspect the source of _wait_for_dependencies
    source = inspect.getsource(DependencyAwareExecutor._wait_for_dependencies)

    # Must reference operation.metadata.get("aggregation_sources")
    # (not operation.parameters.get or parameters["aggregation_sources"])
    assert (
        'metadata.get("aggregation_sources"' in source
        or "metadata.get('aggregation_sources'" in source
    ), (
        "_wait_for_dependencies must read aggregation_sources from operation.metadata, "
        "not operation.parameters"
    )

    # Must NOT read aggregation_sources from parameters
    # (a false positive would be parameters having the key — already caught by test 3)
    assert 'parameters.get("aggregation_sources"' not in source, (
        "_wait_for_dependencies must not read aggregation_sources from operation.parameters"
    )


# ── Test 9: CLI error handler imports cancelled_exc_classes (not get_cancelled_exc_class) ──


def test_error_handler_uses_cached_exc_class():
    """The CLI orchestrate module must detect cancellation with a loop-safe
    helper (is_cancelled or cancelled_exc_classes), not get_cancelled_exc_class.

    get_cancelled_exc_class() (anyio) requires a running event loop and raises
    NoEventLoopError in exception handlers that run after the loop exits.
    cancelled_exc_classes() returns the cached classes; is_cancelled() wraps it.
    """
    import ast
    from pathlib import Path

    orchestrate_init = Path(__file__).parent.parent.parent / "lionagi/cli/orchestrate/__init__.py"
    source = orchestrate_init.read_text()

    tree = ast.parse(source)

    # Collect names imported from any lionagi.ln.concurrency module — the
    # package re-exports the loop-safe helpers from its .errors submodule.
    imported_concurrency: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom):
            module = node.module or ""
            if "ln.concurrency" in module:
                for alias in node.names:
                    imported_concurrency.add(alias.asname or alias.name)

    # The error handler must detect cancellation via a loop-safe helper:
    # is_cancelled() (which wraps the cached classes) or cancelled_exc_classes()
    # directly. Both read the module-level cache and never require a running
    # loop, so they stay safe in handlers that run after the loop exits.
    loop_safe = {"is_cancelled", "cancelled_exc_classes"}
    assert imported_concurrency & loop_safe, (
        "CLI orchestrate __init__ must import a loop-safe cancellation check "
        "(is_cancelled or cancelled_exc_classes), not get_cancelled_exc_class"
    )

    # Confirm the dangerous variant is NOT directly called (it may be imported
    # elsewhere but must not appear as a bare call in this file)
    # We check for the pattern in source as a belt-and-suspenders guard.
    # Allow the import of the name itself for completeness testing, but not a call.
    lines_with_get = [
        (i + 1, line)
        for i, line in enumerate(source.splitlines())
        if "get_cancelled_exc_class()" in line and not line.strip().startswith("#")
    ]
    assert not lines_with_get, (
        f"get_cancelled_exc_class() is called directly in orchestrate/__init__.py at "
        f"lines {lines_with_get} — use cancelled_exc_classes() instead"
    )


# ── Test 10: HookBus lifecycle — default hooks are registered ────────────────


def test_hook_bus_lifecycle_integration():
    """Session.hooks returns a properly-wired bus with all DEFAULT_HOOKS registered.

    This is an end-to-end check of the bus lifecycle:
    1. session.hooks creates the bus
    2. build_session_bus wires DEFAULT_HOOKS
    3. The bus is bound to the observer
    4. Default hook points all have at least one handler
    """
    from lionagi.hooks import DEFAULT_HOOKS, HookBus, HookPoint
    from lionagi.session.session import Session

    session = Session()
    bus = session.hooks

    assert isinstance(bus, HookBus)

    # All four default hook points must have handlers registered
    for point in (
        HookPoint.SESSION_START,
        HookPoint.SESSION_END,
        HookPoint.MESSAGE_ADD,
        HookPoint.BRANCH_CREATE,
    ):
        handlers = bus.handlers_for(point)
        assert len(handlers) >= 1, (
            f"HookPoint.{point.name} must have at least one default handler registered; "
            f"got {handlers!r}"
        )
        # The handlers must come from DEFAULT_HOOKS
        default_handlers = DEFAULT_HOOKS[point]
        for h in default_handlers:
            assert h in handlers, (
                f"Default handler {h.__name__!r} missing from bus for {point.name}"
            )

    # Bus must be bound to the session observer (ADR-0076)
    assert bus._observer is session.observer


def test_hook_bus_handlers_for_returns_shallow_copy():
    """handlers_for() must return a copy so callers cannot mutate the bus state."""
    from lionagi.hooks import HookPoint
    from lionagi.session.session import Session

    session = Session()
    bus = session.hooks

    handlers = bus.handlers_for(HookPoint.SESSION_START)
    original_len = len(handlers)

    # Mutating the returned list must not affect the bus
    handlers.clear()
    assert len(bus.handlers_for(HookPoint.SESSION_START)) == original_len, (
        "handlers_for() must return a shallow copy — mutating it must not affect the bus"
    )


# ── Test 11: aggregation metadata survives graph round-trip ──────────────────


def test_aggregation_metadata_survives_without_parameters_leak():
    """Aggregation node must have exactly the keys expected in metadata and
    exactly 'instruction' in parameters — no cross-contamination.

    This directly verifies the regression: aggregation_sources must NOT appear
    in the parameters dict that gets spread into branch.operate(**params).
    """
    from lionagi.casts.emission import TaskAssignment
    from lionagi.operations.node import Operation
    from lionagi.orchestration.patterns import build_fanout_graph
    from lionagi.session.branch import Branch
    from lionagi.session.session import Session

    session = Session()
    roles = {}
    for name in ("analyst", "researcher", "synthesizer"):
        b = Branch(name=name)
        session.include_branches(b)
        roles[name] = b

    assignments = [
        TaskAssignment(task="Analyse market data", assignee="analyst"),
        TaskAssignment(task="Research competitors", assignee="researcher"),
    ]

    graph, worker_ids = build_fanout_graph(
        session, assignments, roles, synthesis_role="synthesizer"
    )

    synth_node = next(
        (
            n
            for n in graph.internal_nodes.values()
            if isinstance(n, Operation) and n.metadata.get("aggregation")
        ),
        None,
    )
    assert synth_node is not None

    # Metadata checks
    meta = synth_node.metadata
    assert meta["aggregation"] is True
    assert meta["aggregation_count"] == 2
    assert len(meta["aggregation_sources"]) == 2
    # Sources are stored as str(w.id) — compare as strings
    for wid in worker_ids:
        assert str(wid) in meta["aggregation_sources"], (
            f"Worker id {wid!r} missing from aggregation_sources"
        )

    # Parameters checks — only instruction, no aggregation leakage
    params = synth_node.parameters
    assert isinstance(params, dict)
    assert set(params.keys()) == {"instruction"}, (
        f"Synthesis node parameters must only contain 'instruction'; "
        f"got extra keys: {set(params.keys()) - {'instruction'}}"
    )
