from .imports import *

from flask import Response, stream_with_context
from pydantic import BaseModel
from typing import Optional, List


def sse_event(payload: dict) -> bytes:
    return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n".encode("utf-8")


def event_to_sse(ev) -> bytes:
    """Serialize a dispatch StreamEvent to the browser's SSE wire shape.

    token/done/error get their minimal browser payloads; everything else
    (status / provisioning progress / continuation markers — including events
    relayed from a GPU worker) rides through verbatim via model_dump().
    """
    t = getattr(ev, "type", None)
    if t == "token":
        return sse_event({"type": "token", "text": ev.text})
    if t == "done":
        return sse_event({"type": "done", "finish_reason": ev.finish_reason})
    if t == "error":
        return sse_event({"type": "error", "message": ev.message})
    return sse_event(ev.model_dump())


def chat_iter_sync(agen):
    """Drive an async generator from Flask's synchronous WSGI context."""
    loop = asyncio.new_event_loop()

    try:
        asyncio.set_event_loop(loop)

        while True:
            try:
                item = loop.run_until_complete(agen.__anext__())

                if isinstance(item, str):
                    item = item.encode("utf-8")

                yield item

            except StopAsyncIteration:
                break

    finally:
        try:
            loop.run_until_complete(loop.shutdown_asyncgens())
        except Exception:
            pass

        asyncio.set_event_loop(None)
        loop.close()


def _resolve_max_new_tokens(body: ChatBody) -> int:
    """Default to the model's full context when the client didn't cap it.

    A tool, not a service — so when max_new_tokens is omitted we give the model
    as much room as it has. The engine auto-continues past this per-call cap, so
    this is the per-pass budget, not a hard ceiling on total output.
    """
    if body.max_new_tokens:
        return body.max_new_tokens
    try:
        from abstract_hugpy.imports.config.main import get_model_config
        cfg = get_model_config(body.model_key) if body.model_key else None
        ctx = getattr(cfg, "model_max_length", None)
        if ctx and int(ctx) > 0:
            return int(ctx)
    except Exception:
        pass
    # Fall back to the global default cap.
    try:
        from abstract_hugpy.imports.src.constants.constants import DEFAULT_MAX_TOKENS
        return int(DEFAULT_MAX_TOKENS)
    except Exception:
        return 4096


async def stream_events(body: ChatBody):
    """Build prompt_kwargs and stream the unified chat engine to SSE.

    The route is deliberately dumb: it does NOT decide local vs worker. It hands
    prompt_kwargs to execute_chat_stream, which drives resolve() — and resolve()
    is the single place that picks in-process / placement-peer / live-GPU-worker
    and falls back to local. So local and worker chat now stream identically
    (token-by-token, with auto-continuation past the cap), and there is no
    separate worker-offload path in this route anymore.
    """
    from abstract_hugpy.managers.dispatch import execute_chat_stream

    prompt_kwargs = {}
    if body.max_new_tokens:
        # Explicit cap from the client -> honor it (bounded, per-call).
        prompt_kwargs["max_new_tokens"] = body.max_new_tokens
    else:
        # No cap requested -> run unbounded: the runner generates chunk-by-chunk
        # until the model naturally stops, so the response is never truncated by
        # a token limit. (Per-chunk size uses the model's context.)
        prompt_kwargs["unbounded"] = True
        prompt_kwargs["max_new_tokens"] = _resolve_max_new_tokens(body)

    if body.model_key:
        prompt_kwargs["model_key"] = body.model_key

    if body.temperature is not None:
        prompt_kwargs["temperature"] = body.temperature

    if body.do_sample is not None:
        prompt_kwargs["do_sample"] = body.do_sample

    if body.messages:
        prompt_kwargs["messages"] = messages_to_dicts(body.messages)
    else:
        prompt_kwargs["prompt"] = body.prompt

    if body.file:
        prompt_kwargs["file"] = body.file
    if body.images:
        prompt_kwargs["images"] = body.images
    if body.request_id:
        # Stable id the engine threads through every continuation pass; also lets
        # the browser correlate the stream.
        prompt_kwargs["request_id"] = body.request_id

    # Text-only chat to a multi-task (e.g. vision) model: route to its
    # text-generation task instead of the default image-text-to-text, so a
    # plain prompt uses the text runner. The vision runner requires an image
    # and would otherwise fail validation. Only do this when no image is given
    # and the model actually lists text-generation.
    if not body.images and not body.file and body.model_key:
        try:
            from abstract_hugpy.imports.config.main import get_model_config
            cfg = get_model_config(body.model_key)
            tasks = getattr(cfg, "tasks", None) or []
            primary = getattr(cfg, "primary_task", None)
            if primary != "text-generation" and "text-generation" in tasks:
                prompt_kwargs["task"] = "text-generation"
        except Exception:
            pass

    logger.info("prompt_kwargs == %s", prompt_kwargs)

    try:
        async for event in execute_chat_stream(**prompt_kwargs):
            yield event_to_sse(event)
    except Exception as exc:
        logger.exception("stream_events failed")
        yield sse_event({"type": "error", "message": str(exc)})


def chat_stream(mimetype=None, headers=None, **kwargs):
    logger.info(kwargs)
    body = ChatBody(**kwargs)

    return Response(
        stream_with_context(chat_iter_sync(stream_events(body))),
        mimetype=mimetype or "text/event-stream",
        headers=headers or {
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Connection": "keep-alive",
        },
        direct_passthrough=True,
    )
