"""
Copyright (C) 2026 Garudex Labs.  All Rights Reserved.
Caracal, a product of Garudex Labs

SDK primitives: spawn an agent session and delegate authority as async context managers.
"""

from __future__ import annotations

from contextlib import asynccontextmanager
from dataclasses import replace
from collections.abc import AsyncGenerator, Awaitable, Callable

from .context import CaracalContext, current, _ctx_var
from .coordinator import (
    AgentKind,
    CoordinatorClient,
    DelegationConstraints,
    DelegationRequest,
    SpawnRequest,
    create_delegation,
    spawn_agent,
    terminate_agent,
)
from .json_types import JsonObject


LifecycleHook = Callable[[CaracalContext], Awaitable[None]]


@asynccontextmanager
async def spawn(
    *,
    coordinator: CoordinatorClient,
    zone_id: str,
    application_id: str,
    subject_token: str,
    parent_id: str | None = None,
    parent_ctx: CaracalContext | None = None,
    kind: AgentKind = AgentKind.INSTANCE,
    ttl_seconds: int | None = None,
    metadata: JsonObject | None = None,
    trace_id: str | None = None,
    on_agent_start: LifecycleHook | None = None,
    on_agent_end: LifecycleHook | None = None,
) -> AsyncGenerator[CaracalContext, None]:
    """Spawn a child agent session and bind it to the current task.

    ``parent_ctx`` overrides the bound :func:`current` lookup; pass it
    explicitly when the orchestrator owns the parent context but the
    spawn runs on a different task (asyncio TaskGroup, thread pool,
    framework worker) where the parent's contextvar is not visible.
    """
    parent = parent_ctx if parent_ctx is not None else current()
    parent_agent_session_id = parent_id or (parent.agent_session_id if parent else None)
    bearer = subject_token

    res = await spawn_agent(
        coordinator,
        bearer,
        SpawnRequest(
            zone_id=zone_id,
            application_id=application_id,
            parent_id=parent_agent_session_id,
            kind=kind,
            ttl_seconds=ttl_seconds,
            metadata=metadata,
        ),
    )

    ctx = CaracalContext(
        subject_token=bearer,
        zone_id=zone_id,
        client_id=application_id,
        agent_session_id=res.agent_session_id,
        parent_edge_id=parent.delegation_edge_id if parent else None,
        session_id=parent.session_id if parent else None,
        trace_id=trace_id or (parent.trace_id if parent else None),
        hop=parent.hop if parent else 0,
    )

    if on_agent_start is not None:
        await on_agent_start(ctx)

    token = _ctx_var.set(ctx)
    try:
        yield ctx
    finally:
        _ctx_var.reset(token)
        if on_agent_end is not None:
            await on_agent_end(ctx)
        if kind != AgentKind.SERVICE:
            await terminate_agent(coordinator, bearer, zone_id, res.agent_session_id)


@asynccontextmanager
async def delegate(
    *,
    coordinator: CoordinatorClient,
    to_agent_session_id: str,
    to_application_id: str,
    scopes: list[str],
    resource_id: str | None = None,
    constraints: DelegationConstraints | None = None,
    ttl_seconds: int | None = None,
) -> AsyncGenerator[CaracalContext, None]:
    ctx = current()
    if ctx is None or not ctx.agent_session_id:
        raise RuntimeError("delegate requires an active agent session in context")

    res = await create_delegation(
        coordinator,
        ctx.subject_token,
        DelegationRequest(
            zone_id=ctx.zone_id,
            issuer_application_id=ctx.client_id,
            source_session_id=ctx.agent_session_id,
            target_session_id=to_agent_session_id,
            receiver_application_id=to_application_id,
            resource_id=resource_id,
            scopes=scopes,
            constraints=constraints,
            ttl_seconds=ttl_seconds,
        ),
    )

    child = replace(
        ctx,
        parent_edge_id=ctx.delegation_edge_id,
        delegation_edge_id=res.delegation_edge_id,
        hop=ctx.hop + 1,
    )
    token = _ctx_var.set(child)
    try:
        yield child
    finally:
        _ctx_var.reset(token)


@asynccontextmanager
async def delegate_to_spawn(
    *,
    coordinator: CoordinatorClient,
    zone_id: str,
    application_id: str,
    subject_token: str,
    scopes: list[str],
    resource_id: str | None = None,
    parent_ctx: CaracalContext | None = None,
    constraints: DelegationConstraints | None = None,
    delegation_ttl_seconds: int | None = None,
    kind: AgentKind = AgentKind.INSTANCE,
    ttl_seconds: int | None = None,
    metadata: JsonObject | None = None,
    trace_id: str | None = None,
    on_agent_start: LifecycleHook | None = None,
    on_agent_end: LifecycleHook | None = None,
) -> AsyncGenerator[CaracalContext, None]:
    """Atomic spawn + delegate for fan-out workflows.

    The parent context must be active (or supplied via ``parent_ctx``).
    The child session is created and the parent→child delegation edge is
    recorded before the child context is yielded, so callers can safely
    hand the resulting context off to a background task without racing
    the parent's lifecycle.

    ``parent_ctx`` overrides the bound :func:`current` lookup; pass it
    explicitly from background tasks or worker pools that do not inherit
    the orchestrator's contextvar.
    """
    parent = parent_ctx if parent_ctx is not None else current()
    if parent is None or not parent.agent_session_id:
        raise RuntimeError("delegate_to_spawn requires an active agent session in context")

    spawn_res = await spawn_agent(
        coordinator,
        subject_token,
        SpawnRequest(
            zone_id=zone_id,
            application_id=application_id,
            parent_id=parent.agent_session_id,
            kind=kind,
            ttl_seconds=ttl_seconds,
            metadata=metadata,
        ),
    )

    delegation_res = await create_delegation(
        coordinator,
        parent.subject_token,
        DelegationRequest(
            zone_id=parent.zone_id,
            issuer_application_id=parent.client_id,
            source_session_id=parent.agent_session_id,
            target_session_id=spawn_res.agent_session_id,
            receiver_application_id=application_id,
            resource_id=resource_id,
            scopes=scopes,
            constraints=constraints,
            ttl_seconds=delegation_ttl_seconds,
        ),
    )

    ctx = CaracalContext(
        subject_token=subject_token,
        zone_id=zone_id,
        client_id=application_id,
        agent_session_id=spawn_res.agent_session_id,
        delegation_edge_id=delegation_res.delegation_edge_id,
        parent_edge_id=parent.delegation_edge_id,
        session_id=parent.session_id,
        trace_id=trace_id or parent.trace_id,
        hop=parent.hop + 1,
    )

    if on_agent_start is not None:
        await on_agent_start(ctx)

    token = _ctx_var.set(ctx)
    try:
        yield ctx
    finally:
        _ctx_var.reset(token)
        if on_agent_end is not None:
            await on_agent_end(ctx)
        if kind != AgentKind.SERVICE:
            await terminate_agent(coordinator, subject_token, zone_id, spawn_res.agent_session_id)
