# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

"""ds4 (DeepSeek V4) proxy backend.

ds4-server already speaks the OpenAI HTTP API, so this backend is a thin proxy: it
forwards chat/completion requests to the managed ``ds4-server`` subprocess (whose
lifecycle is owned by :mod:`codai.api.ds4_worker`) and adapts the responses to the
:class:`~codai.backends.base.ModelBackend` contract the model manager expects.

Tool/think parsing is handled the same way as the other backends — by
``ModelParserAdapter`` over the returned text — so tools are not forwarded to
ds4-server; the text-level ``DeepSeekParser`` extracts ``<think>`` and tool calls.
"""

import asyncio
import threading
from typing import AsyncGenerator, Dict, List, Optional

from codai.backends.base import ModelBackend


class Ds4Backend(ModelBackend):
    """Proxy backend that routes generation to a managed ds4-server."""

    def __init__(self, cfg=None):
        # cfg is a codai.config.Ds4Config. When omitted, resolve the active one.
        if cfg is None:
            from codai.config import Ds4Config
            cfg = Ds4Config()
        self._cfg = cfg
        self._model_id = getattr(cfg, "model_id", "deepseek-v4") or "deepseek-v4"
        self._url: Optional[str] = None
        self._ctx = int(getattr(cfg, "ctx", 100000) or 100000)
        self._last_usage: Dict = {}

    # ------------------------------------------------------------------ #
    # lifecycle
    # ------------------------------------------------------------------ #
    def load_model(self, model_name: str, **kwargs) -> None:
        from codai.api import ds4_worker
        if model_name:
            self._model_id = model_name
        self._url = ds4_worker.ensure_service(self._cfg)

    def get_model_name(self) -> str:
        return self._model_id

    def get_context_size(self) -> int:
        return self._ctx

    def get_last_usage(self) -> dict:
        return dict(self._last_usage)

    def cleanup(self) -> None:
        from codai.api import ds4_worker
        ds4_worker.stop_service(getattr(self._cfg, "model_id", self._model_id))
        self._url = None

    # ------------------------------------------------------------------ #
    # helpers
    # ------------------------------------------------------------------ #
    def _base(self) -> str:
        if not self._url:
            raise RuntimeError("ds4 service not started")
        return self._url

    def _store_usage(self, usage: dict) -> None:
        if usage:
            self._last_usage = {
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0),
            }

    def format_messages(self, messages) -> str:
        # ds4-server applies DeepSeek V4's own chat template server-side; this is only
        # used by callers that need a flat prompt string.
        parts = []
        for m in messages:
            role = m.get("role") if isinstance(m, dict) else getattr(m, "role", "")
            content = m.get("content") if isinstance(m, dict) else getattr(m, "content", "")
            parts.append(f"{role}: {content}")
        return "\n".join(parts)

    def _chat_payload(self, messages, max_tokens, temperature, top_p, stop, stream):
        payload = {
            "model": self._model_id,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "stream": stream,
        }
        if max_tokens:
            payload["max_tokens"] = max_tokens
        if stop:
            payload["stop"] = stop
        return payload

    # ------------------------------------------------------------------ #
    # chat-level generation (preferred by the manager)
    # ------------------------------------------------------------------ #
    def generate_chat(self, messages: List[Dict], max_tokens=None, temperature=0.7,
                      top_p=1.0, stop=None, tools=None, response_format=None):
        import requests
        payload = self._chat_payload(messages, max_tokens, temperature, top_p, stop, False)
        if response_format and response_format.get("type") == "json_object":
            payload["response_format"] = {"type": "json_object"}
        r = requests.post(self._base() + "/v1/chat/completions", json=payload, timeout=3600)
        r.raise_for_status()
        data = r.json()
        self._store_usage(data.get("usage", {}))
        return data["choices"][0]["message"].get("content") or ""

    async def generate_chat_stream(self, messages: List[Dict], max_tokens=None,
                                   temperature=0.7, top_p=1.0, stop=None, tools=None,
                                   response_format=None) -> AsyncGenerator[str, None]:
        payload = self._chat_payload(messages, max_tokens, temperature, top_p, stop, True)
        async for chunk in self._stream(self._base() + "/v1/chat/completions", payload,
                                        delta_key="delta"):
            yield chunk

    # ------------------------------------------------------------------ #
    # plain completion (fallback path)
    # ------------------------------------------------------------------ #
    def generate(self, prompt: str, max_tokens=None, temperature: float = 0.7,
                 top_p: float = 1.0, stop=None, repeat_penalty: float = 1.0,
                 presence_penalty: float = 0.0, frequency_penalty: float = 0.0) -> str:
        return self.generate_chat([{"role": "user", "content": prompt}],
                                  max_tokens, temperature, top_p, stop)

    async def generate_stream(self, prompt: str, max_tokens=None, temperature: float = 0.7,
                              top_p: float = 1.0, stop=None, repeat_penalty: float = 1.0,
                              presence_penalty: float = 0.0,
                              frequency_penalty: float = 0.0) -> AsyncGenerator[str, None]:
        async for chunk in self.generate_chat_stream(
                [{"role": "user", "content": prompt}], max_tokens, temperature, top_p, stop):
            yield chunk

    # ------------------------------------------------------------------ #
    # SSE streaming: iterate the blocking requests stream on a worker thread
    # and hand chunks to the event loop through an asyncio.Queue.
    # ------------------------------------------------------------------ #
    async def _stream(self, url: str, payload: dict, delta_key: str
                      ) -> AsyncGenerator[str, None]:
        import json
        loop = asyncio.get_event_loop()
        queue: asyncio.Queue = asyncio.Queue()
        _SENTINEL = object()

        def _worker():
            import requests
            try:
                with requests.post(url, json=payload, stream=True, timeout=3600) as r:
                    r.raise_for_status()
                    for raw in r.iter_lines(decode_unicode=True):
                        if not raw or not raw.startswith("data:"):
                            continue
                        data = raw[len("data:"):].strip()
                        if data == "[DONE]":
                            break
                        try:
                            obj = json.loads(data)
                        except ValueError:
                            continue
                        choice = (obj.get("choices") or [{}])[0]
                        text = (choice.get(delta_key) or {}).get("content") or ""
                        if text:
                            loop.call_soon_threadsafe(queue.put_nowait, text)
                        if obj.get("usage"):
                            self._store_usage(obj["usage"])
                        if choice.get("finish_reason"):
                            break
            except Exception as exc:  # surface to the consumer
                loop.call_soon_threadsafe(queue.put_nowait, exc)
            finally:
                loop.call_soon_threadsafe(queue.put_nowait, _SENTINEL)

        threading.Thread(target=_worker, daemon=True).start()
        while True:
            item = await queue.get()
            if item is _SENTINEL:
                break
            if isinstance(item, Exception):
                raise item
            yield item
