from __future__ import annotations

import asyncio
import json
import sys
import traceback
from collections.abc import Awaitable, Callable
from typing import Any

from fastapi.responses import StreamingResponse
from matrx_utils import vcprint

from matrx_service.context.app_context import AppContext
from matrx_service.context.emitter_protocol import Emitter


def create_streaming_response(
    ctx: AppContext,
    task_fn: Callable[..., Awaitable[None]],
    *task_args: Any,
    initial_message: str = "Connecting...",
    debug_label: str = "stream",
) -> StreamingResponse:
    """Wire a streaming task to an NDJSON response.

    Context is already fully set by AuthMiddleware (both request.state and
    the ContextVar). Callers must call ctx.extend() with any route-specific
    fields (conversation_id, debug, etc.) BEFORE calling this function.

    Parameters
    ----------
    ctx             AppContext from context_dep — already live in the ContextVar.
    task_fn         async def _run_something(emitter: StreamEmitter, *task_args)
    *task_args      Forwarded to task_fn after emitter.
    initial_message First status_update event sent to the client.
    debug_label     Used in error log messages.
    """
    if ctx.emitter is None:
        raise RuntimeError("AppContext.emitter must be set before create_streaming_response()")
    emitter: Emitter = ctx.emitter

    async def _stream():
        yield json.dumps({
            "event": "status_update",
            "data": {
                "status": "connected",
                "system_message": "Stream established",
                "user_message": initial_message,
            },
        }) + "\n"

        if ctx.conversation_id:
            yield json.dumps({
                "event": "data",
                "data": {
                    "event": "conversation_id",
                    "conversation_id": ctx.conversation_id,
                },
            }) + "\n"

        async def _run_with_error_handling():
            try:
                await task_fn(emitter, *task_args)
            except asyncio.CancelledError:
                print(
                    f"\n[{debug_label}] Task cancelled (client disconnected)",
                    file=sys.stderr, flush=True,
                )
                vcprint("Task cancelled (client disconnected)", f"[{debug_label}]", color="yellow")
            except Exception as e:
                tb = traceback.format_exc()
                # Write directly to stderr — bypasses all logging infrastructure,
                # never silenced by uvicorn log config or level filters.
                print(
                    f"\n[{debug_label}] TASK CRASHED — {type(e).__name__}: {e}\n{tb}",
                    file=sys.stderr, flush=True,
                )
                vcprint(e, f"[{debug_label}] TASK CRASHED — {type(e).__name__}", color="red")
                await emitter.fatal_error(
                    error_type=f"{debug_label.lower()}_error",
                    message=str(e),
                    user_message=f"{debug_label} failed.",
                )

        task = asyncio.create_task(_run_with_error_handling())
        emitter.set_task(task)

        async for item in emitter.generate():
            yield item

        # After the generator exits: if the task raised an unhandled exception
        # (shouldn't happen — _run_with_error_handling catches everything — but
        # belt-and-suspenders: log it so it's never silent).
        if task.done() and not task.cancelled():
            exc = task.exception()
            if exc is not None:
                tb = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
                print(
                    f"\n[{debug_label}] TASK EXCEPTION ESCAPED HANDLER — "
                    f"{type(exc).__name__}: {exc}\n{tb}",
                    file=sys.stderr, flush=True,
                )

    headers: dict[str, str] = {}
    if ctx.request_id:
        headers["X-Request-ID"] = ctx.request_id
    if ctx.conversation_id:
        headers["X-Conversation-ID"] = ctx.conversation_id

    return StreamingResponse(
        _stream(),
        media_type="application/x-ndjson",
        headers=headers,
    )
