from .imports import *
# base_runner.py
class LlamaCppBaseRunner(ABC):
    """Shared scaffolding — event loop, unbounded loop, logging.
    Subclasses implement only the raw I/O."""

    model_key: str

    # --- abstract I/O hooks ------------------------------------------------

    @abstractmethod
    async def _iter_stream(
        self,
        messages: list[dict],
        max_tokens: int,
        temp: float,
        top_p: float,
    ) -> AsyncIterator[tuple[str, Optional[str]]]:
        """Yield (text_chunk, finish_reason_or_None) pairs from the backend."""
        ...
    @abstractmethod
    def _chat_complete(
        self,
        messages: list[dict],
        max_tokens: int,
        temp: float,
        top_p: float,
        stop: Optional[list[str]],
    ) -> tuple[str, str]:
        """Chat-template path. Return (text, finish_reason)."""
        ...

    @abstractmethod
    def _raw_complete(
        self,
        prompt: str,
        max_tokens: int,
        temp: float,
        top_p: float,
        stop: Optional[list[str]],
        return_full_text: bool,
    ) -> tuple[str, str]:
        """Raw-prompt fallback path. Return (text, finish_reason)."""
        ...
    # _blocking_complete is now FINAL — no override needed in subclasses
    def _blocking_complete(
        self,
        messages: list[dict] | str,
        max_tokens: int,
        temp: float,
        top_p: float,
        stop: Optional[list[str]],
        use_chat_template: bool,
        return_full_text: bool,
    ) -> tuple[str, str]:
        if use_chat_template and isinstance(messages, list):
            return self._chat_complete(messages, max_tokens, temp, top_p, stop)

        prompt = (
            messages
            if isinstance(messages, str)
            else messages_to_prompt_from_dicts(messages)
        )
        return self._raw_complete(prompt, max_tokens, temp, top_p, stop, return_full_text)
    # --- multimodal: fold an image attachment into the chat ----------------

    def _attach_image(self, messages: list, req: ChatRequest) -> list:
        """Vision GGUFs only: fold image attachments into the latest user turn as
        OpenAI ``image_url`` parts so the multimodal chat handler sees them.

        Two sources, both supported: ``req.images`` — inline base64 (raw or a full
        ``data:`` URI), the no-upload path used by the chat box — and ``req.file``,
        an uploaded image path. No-op for text models / non-image files / no
        attachments, so text chat and non-vision runners are completely unaffected.
        """
        if not getattr(self, "is_vision", False):
            return messages
        import base64, mimetypes
        parts: list = []
        # Inline base64 images (no upload round-trip).
        for img in (getattr(req, "images", None) or []):
            if not img:
                continue
            url = img if str(img).startswith("data:") else f"data:image/png;base64,{img}"
            parts.append({"type": "image_url", "image_url": {"url": url}})
        # Uploaded image file path.
        path = getattr(req, "file", None)
        if path and os.path.isfile(path):
            mime = mimetypes.guess_type(path)[0] or "image/png"
            if mime.startswith("image/"):
                with open(path, "rb") as fh:
                    data_uri = f"data:{mime};base64," + base64.b64encode(fh.read()).decode("ascii")
                parts.append({"type": "image_url", "image_url": {"url": data_uri}})
        if not parts:
            return messages
        for m in reversed(messages):
            if m.get("role", "user") == "user":
                content = m.get("content")
                if isinstance(content, list):
                    content.extend(parts)
                elif content:
                    m["content"] = [{"type": "text", "text": content}, *parts]
                else:
                    m["content"] = list(parts)
                break
        else:
            messages.append({"role": "user", "content": list(parts)})
        return messages

    # --- shared streaming --------------------------------------------------

    async def stream_chat(
        self,
        req: ChatRequest,
        cancel_event: Optional[asyncio.Event] = None,
    ) -> AsyncIterator[StreamEvent]:
        max_tokens = resolve_max_tokens(req.max_new_tokens)
        temp      = resolve_temperature(req.temperature, req.do_sample)
        top_p     = resolve_top_p(req.top_p)
        messages  = self._attach_image(messages_to_dicts(req.messages), req)
        output_chunks = 0
        last_finish: Optional[str] = None

        try:
            async for text, fr in self._iter_stream(messages, max_tokens, temp, top_p):
                if cancel_event and cancel_event.is_set():
                    self._log_done(req, "cancelled", output_chunks, max_tokens)
                    yield DoneEvent(request_id=req.request_id, input_tokens=0,
                                   output_chunks=output_chunks, finish_reason="cancelled")
                    return
                if text:
                    output_chunks += 1
                    yield TokenEvent(request_id=req.request_id, text=text)
                if fr is not None:
                    last_finish = fr

            mapped = map_finish_reason(last_finish)
            self._log_done(req, mapped, output_chunks, max_tokens)
            yield DoneEvent(request_id=req.request_id, input_tokens=0,
                           output_chunks=output_chunks, finish_reason=mapped)
        except Exception as exc:
            logger.exception("stream_chat failed: model=%s req=%s", self.model_key, req.request_id)
            yield ErrorEvent(request_id=req.request_id, message=f"{type(exc).__name__}: {exc}")

    # --- shared unbounded streaming ----------------------------------------

    async def stream_chat_unbounded(
        self,
        req: ChatRequest,
        cancel_event: Optional[asyncio.Event] = None,
        *,
        chunk_tokens: int = 1024,
        max_chunks: int = None,
    ) -> AsyncIterator[StreamEvent]:
        import os as _os
        if max_chunks is None:
            # High ceiling so "unbounded" really means until-the-model-stops,
            # not a hidden 8-chunk cap. Still bounded so a looping model can't
            # run forever. Override with HUGPY_MAX_CHUNKS.
            try:
                max_chunks = int(_os.environ.get("HUGPY_MAX_CHUNKS", "256"))
            except ValueError:
                max_chunks = 256
        temp     = resolve_temperature(req.temperature, req.do_sample)
        top_p    = resolve_top_p(req.top_p)
        convo    = self._attach_image(messages_to_dicts(req.messages), req)
        output_chunks = 0
        last_finish = "stop"

        try:
            for _ in range(max_chunks):
                if cancel_event and cancel_event.is_set():
                    self._log_done(req, "cancelled", output_chunks, chunk_tokens)
                    yield DoneEvent(request_id=req.request_id, input_tokens=0,
                                   output_chunks=output_chunks, finish_reason="cancelled")
                    return

                piece_text = ""
                chunk_finish: Optional[str] = None

                async for text, fr in self._iter_stream(convo, chunk_tokens, temp, top_p):
                    if cancel_event and cancel_event.is_set():
                        self._log_done(req, "cancelled", output_chunks, chunk_tokens)
                        yield DoneEvent(request_id=req.request_id, input_tokens=0,
                                       output_chunks=output_chunks, finish_reason="cancelled")
                        return
                    if text:
                        output_chunks += 1
                        piece_text += text
                        yield TokenEvent(request_id=req.request_id, text=text)
                    if fr is not None:
                        chunk_finish = fr

                last_finish = chunk_finish or "stop"
                if last_finish != "length" or not piece_text:
                    break

                convo.append({"role": "assistant", "content": piece_text})
                convo.append({"role": "user", "content": "continue"})

            mapped = map_finish_reason(last_finish)
            self._log_done(req, mapped, output_chunks, chunk_tokens)
            yield DoneEvent(request_id=req.request_id, input_tokens=0,
                           output_chunks=output_chunks, finish_reason=mapped)
        except Exception as exc:
            logger.exception("stream_chat_unbounded failed: model=%s req=%s", self.model_key, req.request_id)
            yield ErrorEvent(request_id=req.request_id, message=f"{type(exc).__name__}: {exc}")

    # --- shared non-streaming ----------------------------------------------

    async def generate_text_async(self, messages, **kw) -> str:
        return await asyncio.to_thread(self.generate_text, messages, **kw)

    def generate_text(self, messages, *, max_new_tokens=0, temperature=0.0,
                      top_p=1.0, do_sample=False, use_chat_template=True,
                      return_full_text=False, stop=None, **_) -> str:
        max_tokens = resolve_max_tokens(max_new_tokens)
        temp       = resolve_temperature(temperature, do_sample)
        top_p_val  = resolve_top_p(top_p)
        # _blocking_complete returns (text, finish_reason) per the base contract;
        # unpack so a one-shot run() yields a str (not a tuple) into ChatResult.text.
        text, _finish = self._blocking_complete(
            messages, max_tokens, temp, top_p_val, stop, use_chat_template, return_full_text
        )
        return text

    def generate_text_unbounded(self, messages, *, chunk_tokens=1024,
                                max_chunks=None, temperature=0.0, top_p=1.0,
                                do_sample=False, stop=None, **_) -> str:
        import os as _os
        if max_chunks is None:
            try:
                max_chunks = int(_os.environ.get("HUGPY_MAX_CHUNKS", "256"))
            except ValueError:
                max_chunks = 256
        temp      = resolve_temperature(temperature, do_sample)
        top_p_val = resolve_top_p(top_p)
        accumulated = ""
        convo = list(messages)

        for chunk_idx in range(max_chunks):
            text, finish = self._blocking_complete(
                convo, chunk_tokens, temp, top_p_val, stop,
                use_chat_template=True, return_full_text=False
            )
            accumulated += text
            logger.info("generate_text_unbounded chunk=%s model=%s finish=%s",
                       chunk_idx, self.model_key, finish)
            if finish != "length" or not text:
                break
            convo = convo + [{"role": "assistant", "content": text},
                             {"role": "user", "content": "continue"}]

        return accumulated

    # --- shared internals --------------------------------------------------

    def _log_done(self, req: ChatRequest, finish: str, chunks: int, cap: int) -> None:
        logger.info("stream_chat done: model=%s req=%s finish=%s chunks=%s cap=%s",
                   self.model_key, req.request_id, finish, chunks, cap)
