"""
Unified streaming functions — mirrors packages/ai/src/stream.ts

Provides stream(), complete(), stream_simple(), complete_simple().
"""
from __future__ import annotations

from typing import AsyncGenerator

from .api_registry import get_api_provider
from .env_api_keys import resolve_provider_key
from .types import (
    AssistantMessage,
    AssistantMessageEvent,
    Context,
    EventDone,
    EventError,
    Model,
    SimpleStreamOptions,
    StreamOptions,
)

# Avoid eager `from .register import register_builtins` at module top:
# `register.py` imports `openprogram.providers`, which imports this
# module — fine on the main thread (Python resolves the partial module),
# but a concurrent worker thread that triggers the same import races
# the partial state and raises ImportError. Importing lazily breaks
# the cycle.
_builtins_registered = False


def _ensure_builtins() -> None:
    global _builtins_registered
    if _builtins_registered:
        return
    from .register import register_builtins
    register_builtins()
    _builtins_registered = True


# NOTE: _ensure_builtins() is called once from providers/__init__.py
# at the very end of package init, AFTER stream.py has finished
# loading. Calling it here would re-enter providers (via register.py)
# while __init__.py is mid-way and trip Python's import machinery in
# multi-threaded contexts.


async def stream_simple(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AsyncGenerator[AssistantMessageEvent, None]:
    """
    Stream a response with unified reasoning options.
    Automatically resolves API key from environment if not provided.
    Mirrors streamSimple() from TypeScript.
    """
    opts = options or SimpleStreamOptions()

    # Auto-resolve API key if not set. resolve_provider_key reads the
    # AuthStore (the single key source — no env vars, no config.json).
    if not opts.api_key:
        opts = opts.model_copy(update={"api_key": resolve_provider_key(model.provider)})

    # NOTE: the claude-code Meridian-profile header (x-meridian-profile) is
    # injected one layer down, in openai_completions.stream_simple — that's
    # the single chokepoint EVERY claude-code request passes through. This
    # wrapper is bypassed by some callers (e.g. memory/llm_bridge.py calls
    # the raw api-provider directly), so injecting here would miss them.
    # See docs/design/claude-code-meridian-profile.md.

    provider = get_api_provider(model.api)
    if provider is None:
        raise ValueError(f"No stream function registered for API: {model.api!r}")

    recorded = False
    async for event in provider.stream_simple(model, context, opts):
        # Record AT the terminal event, not after the loop: the consumer
        # (agent_loop) returns the moment it sees the done/error event,
        # leaving this generator suspended at ``yield`` — a post-loop line
        # would never run. The terminal event carries the final message, so
        # we have everything we need to record before yielding it onward.
        if not recorded:
            final = _extract_final(event)
            if final is not None:
                recorded = True
                _record_usage(model, final, opts)
        yield event


async def complete_simple(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessage:
    """
    Get a complete (non-streaming) response.
    Mirrors completeSimple() from TypeScript.
    """
    final_message: AssistantMessage | None = None

    async for event in stream_simple(model, context, options):
        final_message = _extract_final(event) or final_message

    if final_message is None:
        raise RuntimeError("Stream completed without a final message")

    return final_message


async def stream(
    model: Model,
    context: Context,
    options: StreamOptions | None = None,
) -> AsyncGenerator[AssistantMessageEvent, None]:
    """
    Stream with provider-specific options (no reasoning normalization).
    Mirrors stream() from TypeScript.
    """
    opts = options or StreamOptions()

    # Auto-resolve API key from the AuthStore if not set (same as stream_simple)
    if not opts.api_key:
        opts = opts.model_copy(update={"api_key": resolve_provider_key(model.provider)})

    provider = get_api_provider(model.api)
    if provider is None:
        raise ValueError(f"No stream function registered for API: {model.api!r}")

    recorded = False
    async for event in provider.stream(model, context, opts):
        if not recorded:
            final = _extract_final(event)
            if final is not None:
                recorded = True
                _record_usage(model, final, opts)
        yield event


async def complete(
    model: Model,
    context: Context,
    options: StreamOptions | None = None,
) -> AssistantMessage:
    """
    Get a complete response with provider-specific options.
    Mirrors complete() from TypeScript.
    """
    final_message: AssistantMessage | None = None

    async for event in stream(model, context, options):
        final_message = _extract_final(event) or final_message

    if final_message is None:
        raise RuntimeError("Stream completed without a final message")

    return final_message


def _record_usage(model: Model, final, options) -> None:
    """Single metering chokepoint: every stream()/stream_simple() ends here,
    so every LLM call this module serves gets recorded. Best-effort — a
    metering failure must never surface to the caller. The call source
    (chat / compaction / memory / …) is read from the contextvar set by
    the caller's ``usage_scope(...)``; session_id comes off the options.
    """
    if final is None:
        return
    try:
        from openprogram.metering import record_message
        session_id = getattr(options, "session_id", None) if options else None
        record_message(model, final, session_id=session_id)
    except Exception:
        pass


def _extract_final(event) -> AssistantMessage | None:
    """
    Pull the AssistantMessage out of a terminal event.
    Provider implementations sometimes yield dicts rather than BaseModel
    instances — normalize both shapes.
    """
    etype = event["type"] if isinstance(event, dict) else getattr(event, "type", None)
    if etype == "done":
        payload = event["message"] if isinstance(event, dict) else event.message
    elif etype == "error":
        payload = event["error"] if isinstance(event, dict) else event.error
    else:
        return None

    if isinstance(payload, AssistantMessage):
        return payload
    if isinstance(payload, dict):
        return AssistantMessage.model_validate(payload)
    return payload
