import pydantic
import structlog
from openai import AsyncOpenAI, PermissionDeniedError

from talemate.client.base import ClientBase, ErrorAction, CommonDefaults
from talemate.client.registry import register
from talemate.emit import emit
from talemate.util import count_tokens

__all__ = [
    "DeepSeekClient",
]
log = structlog.get_logger("talemate")

BASE_URL = "https://api.deepseek.com"

# Edit this to add new models / remove old models
SUPPORTED_MODELS = [
    "deepseek-chat",
    "deepseek-reasoner",
]

JSON_OBJECT_RESPONSE_MODELS = [
    "deepseek-chat",
]


class Defaults(CommonDefaults, pydantic.BaseModel):
    max_token_length: int = 16384
    model: str = "deepseek-chat"


@register()
class DeepSeekClient(ClientBase):
    """
    DeepSeek client for generating text.
    """

    client_type = "deepseek"
    conversation_retries = 0
    # TODO: make this configurable?
    decensor_enabled = False

    class Meta(ClientBase.Meta):
        name_prefix: str = "DeepSeek"
        title: str = "DeepSeek"
        manual_model: bool = True
        manual_model_choices: list[str] = SUPPORTED_MODELS
        requires_prompt_template: bool = False
        defaults: Defaults = Defaults()
        unified_api_key_config_path: str = "deepseek.api_key"

    @property
    def deepseek_api_key(self):
        return self.config.deepseek.api_key

    @property
    def supported_parameters(self):
        return [
            "temperature",
            "top_p",
            "presence_penalty",
            "max_tokens",
        ]

    def emit_status(self, processing: bool = None):
        error_action = None
        error_message = None
        if processing is not None:
            self.processing = processing

        if self.deepseek_api_key:
            status = "busy" if self.processing else "idle"
        else:
            status = "error"
            error_message = "No API key set"
            error_action = ErrorAction(
                title="Set API Key",
                action_name="openAppConfig",
                icon="mdi-key-variant",
                arguments=[
                    "application",
                    "deepseek_api",
                ],
            )

        if not self.model_name:
            status = "error"
            error_message = "No model loaded"

        self.current_status = status

        data = {
            "error_action": error_action.model_dump() if error_action else None,
            "meta": self.Meta().model_dump(),
            "enabled": self.enabled,
            "error_message": error_message,
        }
        data.update(self._common_status_data())
        emit(
            "client_status",
            message=self.client_type,
            id=self.name,
            details=self.model_name,
            status=status if self.enabled else "disabled",
            data=data,
        )

    def count_tokens(self, content: str):
        if not self.model_name:
            return 0
        return count_tokens(content)

    async def status(self):
        self.emit_status()

    def response_tokens(self, response: str):
        # Count tokens in a response string using the util.count_tokens helper
        return self.count_tokens(response)

    def prompt_tokens(self, prompt: str):
        # Count tokens in a prompt string using the util.count_tokens helper
        return self.count_tokens(prompt)

    async def generate(self, prompt: str, parameters: dict, kind: str):
        """
        Generates text from the given prompt and parameters.
        """

        if not self.deepseek_api_key:
            raise Exception("No DeepSeek API key set")

        client = AsyncOpenAI(api_key=self.deepseek_api_key, base_url=BASE_URL)

        human_message = {"role": "user", "content": prompt.strip()}
        system_message = {"role": "system", "content": self.get_system_message(kind)}

        self.log.debug(
            "generate",
            prompt=prompt[:128] + " ...",
            parameters=parameters,
            system_message=system_message,
        )

        try:
            # Use streaming so we can update_Request_tokens incrementally
            stream = await client.chat.completions.create(
                model=self.model_name,
                messages=[system_message, human_message],
                stream=True,
                **parameters,
            )

            response = ""

            # Iterate over streamed chunks
            async for chunk in stream:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if delta and getattr(delta, "content", None):
                    content_piece = delta.content
                    response += content_piece
                    # Incrementally track token usage
                    self.update_request_tokens(self.count_tokens(content_piece))

            # Save token accounting for whole request
            self._returned_prompt_tokens = self.prompt_tokens(prompt)
            self._returned_response_tokens = self.response_tokens(response)

            return response
        except PermissionDeniedError as e:
            self.log.error("generate error", e=e)
            emit("status", message="DeepSeek API: Permission Denied", status="error")
            return ""
        except Exception:
            raise
