from typing import Optional, Literal

from openkosmos_core.web.model import RestClientConfig
from pydantic import BaseModel

from openkosmos_ai.common.model import AIBaseConfig


class ApiClientConfig(RestClientConfig):
    pass


class DatabaseClientConfig(AIBaseConfig):
    url: str
    create_table: bool = False


class DifyConsoleAuthConfig(BaseModel):
    api_prefix: str = "/v1"
    email: str
    password: str


class DifyClientConfig(ApiClientConfig):
    console_auth: DifyConsoleAuthConfig


class GitRepoConfig(AIBaseConfig):
    repo_dir: str
    remote_branch: str
    remote_url: str
    username: str
    email: str


class ModelGatewayConfig(ApiClientConfig):
    rerank_model: Optional[str] = "bge-reranker-v2-m3"
    embedding_model: Optional[str] = "bge-m3"
    api_prefix: str = "/v1"


class LitellmClientConfig(ApiClientConfig):
    rerank_model: Optional[str] = "bge-reranker-v2-m3"
    embedding_model: Optional[str] = "bge-m3"
    api_prefix: str = "/v1"


class OpenAIClientConfig(AIBaseConfig):
    base_url: str
    api_key: str
    model: Optional[str] = None
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.8
    presence_penalty: Optional[float] = 1.03
    streaming: Optional[bool] = True
    stream_usage: Optional[bool] = True
    extra_body: Optional[dict] = {}


class PlaceholderConfig(AIBaseConfig):
    name: str


class PageParams(BaseModel):
    page: int = 1
    page_size: int = 9999
    orderby: Literal["create_time", "update_time"] = None
    desc: bool = None

    def to_http_params(self, params: dict = {}):
        return self.model_dump(exclude_none=True) | params


class RAGFlowClientConfig(ApiClientConfig):
    api_prefix: str = "/api/v1"


class AgentConfig(AIBaseConfig):
    openai: OpenAIClientConfig


class SkillConfig(AIBaseConfig):
    openai: OpenAIClientConfig


class EmbeddingConfig(BaseModel):
    model_name: str
    model_kwargs: Optional[dict] = {"device": "cpu"}


class VectorStoreConfig(AIBaseConfig):
    embedding: EmbeddingConfig
    index_dir: Optional[str] = None
    index_name: Optional[str] = None
    save_index: bool = False
    load_index: bool = False

    def enable_save(self, enable=True):
        self.save_index = enable
        return self

    def enable_load(self, enable=True):
        self.load_index = enable
        return self


class XinferenceClientConfig(ApiClientConfig):
    rerank_model: Optional[str] = "bge-reranker-v2-m3"
    embedding_model: Optional[str] = "bge-m3"
    api_prefix: str = "/v1"
