from datetime import datetime

from pydantic.main import BaseModel
from typing import Optional, Literal, Any, Dict, List


class Error(BaseModel):
    id: int
    description: str

    def __str__(self):
        error = {
            "id": self.id,
            "description": self.description
        }
        return str(error)


class AssistantRevisionMetadata(BaseModel):
    """
    {
      "key": "string",
      "type": "string",
      "value": "string"
    }
    """
    key: str
    type: str
    value: str

    def __str__(self):
        metadata = {
            "key": self.key,
            "type": self.type,
            "value": self.value
        }
        return str(metadata)


class AssistantRevision(BaseModel):
    """
    {
      "metadata": [
        ...
      ],
      "modelId": "string",
      "modelName": "string",
      "prompt": "string",
      "providerName": "string",
      "revisionDescription": "string",
      "revisionId": "string",
      "revisionName": "string",
      "timestamp": "timestamp"
    }
    """
    metadata: Optional[list[AssistantRevisionMetadata]] = []
    model_id: Optional[str] = None
    model_name: Optional[str] = None
    prompt: Optional[str] = None
    provider_name: Optional[str] = None
    revision_description: Optional[str] = None
    revision_id: int = None
    revision_name: str
    timestamp: Optional[datetime] = None

    def __str__(self):
        revision = {
            "modelId": self.model_id,
            "modelName": self.model_name,
            "prompt": self.prompt,
            "providerName": self.provider_name,
            "revisionDescription": self.revision_description,
            "revisionId": self.revision_id,
            "revisionName": self.revision_name,
            "timestamp": self.timestamp
        }
        if any(self.metadata):
            revision["metadata"] = self.metadata

        return str(revision)


class AssistantIntent(BaseModel):
    """
    {
          "assistantIntentDefaultRevision": "number",
          "assistantIntentDescription": "string",
          "assistantIntentId": "string",
          "assistantIntentName": "string",
          "revisions": [
            ...
          ]
        }
    """
    default_revision: float
    description: str
    id: str
    name: str
    revisions: Optional[list[AssistantRevision]] = []

    def __str__(self):
        intent = {
            "assistantIntentDefaultRevision": self.default_revision,
            "assistantIntentDescription": self.description,
            "assistantIntentId": self.id,
            "assistantIntentName": self.name
        }
        if any(self.revisions):
            intent["revisions"] = self.revisions

        return str(intent)


class Organization(BaseModel):
    id: str
    name: str

    def __str__(self):
        organization = {
            "organizationId": self.id,
            "organizationName": self.name
        }
        return str(organization)


class ProjectSearchProfile(BaseModel):
    """
     {
      "name": "string",
      "description": "string"
    }
    """
    name: str
    description: str

    def __str__(self):
        search_profile = {
            "name": self.name,
            "description": self.description
        }
        return str(search_profile)


class ProjectToken(BaseModel):
    """
     {
      "description": "string",
      "id": "string",
      "name": "string",
      "status": "string", /* Active, Blocked */
      "timestamp": "timestamp"
    }
    """
    description: Optional[str] = None
    token_id: str
    name: str
    status: str
    timestamp: datetime

    def __str__(self):
        token = {
            "description": self.description,
            "id": self.token_id,
            "name": self.name,
            "status": self.status,
            "timestamp": self.timestamp
        }
        return str(token)


class ProjectUsageLimit(BaseModel):
    """
    "hardLimit": "number",                // Upper usage limit
    "id": "string",                       // Usage limit ID
    "relatedEntityName": "string",        // Name of the related entity
    "remainingUsage": "number",           // Remaining usage
    "renewalStatus": "string",            // Renewal status (Renewable, NonRenewable)
    "softLimit": "number",                // Lower usage limit
    "status": "integer",                  // Status (1: Active, 2: Expired, 3: Empty, 4: Cancelled)
    "subscriptionType": "string",         // Subscription type (Freemium, Daily, Weekly, Monthly)
    "usageUnit": "string",                // Usage unit (Requests, Cost)
    "usedAmount": "number",               // Amount used (decimal or scientific notation)
    "validFrom": "timestamp",             // Start date of the usage limit
    "validUntil": "timestamp"             // Expiration or renewal date
    """
    hard_limit: Optional[float] = None
    usage_limit_id: Optional[str] = None
    related_entity_name: Optional[str] = None
    remaining_usage: Optional[float] = None
    renewal_status: Optional[Literal["Renewable", "NonRenewable"]] = None
    soft_limit: Optional[float] = None
    status: Optional[int] = None
    subscription_type: Optional[Literal["Freemium", "Daily", "Weekly", "Monthly"]] = None
    usage_unit: Optional[Literal["Requests", "Cost"]] = None
    used_amount: Optional[float] = None
    valid_from: Optional[datetime] = None
    valid_until: Optional[datetime] = None

    def __str__(self):
        usage_limit = {
            "hardLimit": self.hard_limit,
            "id": self.usage_limit_id,
            "relatedEntityName": self.related_entity_name,
            "remainingUsage": self.remaining_usage,
            "renewalStatus": self.renewal_status,
            "softLimit": self.soft_limit,
            "status": self.status,
            "subscriptionType": self.subscription_type,
            "usageUnit": self.usage_unit,
            "usedAmount": self.used_amount,
            "validFrom": self.valid_from,
            "validUntil": self.valid_until
        }
        return str(usage_limit)


class Project(BaseModel):
    """
     {
      "projectActive": "boolean",
      "projectDescription": "string",
      "projectId": "string",
      "projectName": "string",
      "projectStatus": "integer", /* 0:Active, 2:Hidden */
    }
    """
    organization: Optional[Organization] = None
    active: Optional[bool] = None
    description: Optional[str] = None
    id: Optional[str] = None
    name: Optional[str] = None
    email: Optional[str] = None
    status: Optional[int] = None
    search_profiles: Optional[list[ProjectSearchProfile]] = []
    tokens: Optional[list[ProjectToken]] = []
    usage_limit: Optional[ProjectUsageLimit] = None

    def __str__(self):
        project = {
            "projectName": self.name
        }

        if self.id:
            project["projectId"] = self.id

        if self.organization:
            project["organization"] = self.organization

        if self.active is not None:
            project["projectActive"] = self.active

        if self.description is not None:
            project["projectDescription"] = self.description

        if self.status is not None:
            project["projectStatus"] = self.status

        if self.search_profiles is not None and any(self.search_profiles):
            project["searchProfiles"] = self.search_profiles

        if self.tokens is not None and any(self.tokens):
            project["tokens"] = self.tokens

        if self.usage_limit is not None:
            project["usageLimit"] = self.usage_limit

        return str(project)


class ProjectItem(BaseModel):
    """
    {
      "assistant": "string",
      "intent": "string",
      "timestamp": "string",
      "prompt": "string",
      "output": "string",
      "inputText": "string",
      "status": "string"
    }
    """
    assistant: str
    intent: Optional[str] = None
    timestamp: str
    prompt: Optional[str] = None
    output: Optional[str] = None
    input_text: Optional[str] = None
    status: str

    def __str__(self):
        item = {
            "assistant": self.assistant,
            "intent": self.intent,
            "timestamp": self.timestamp,
            "prompt": self.prompt,
            "output": self.output,
            "inputText": self.input_text,
            "status": self.status
        }
        return str(item)


class GuardrailSettings(BaseModel):
    llm_output: Optional[bool] = False
    input_moderation: Optional[bool] = False
    prompt_injection: Optional[bool] = False

    def to_dict(self):
        return {
            "llmOutputGuardrail": self.llm_output,
            "inputModerationGuardrail": self.input_moderation,
            "promptInjectionGuardrail": self.prompt_injection
        }

    def __str__(self):
        settings = self.to_dict()
        return str(settings)


class LlmSettings(BaseModel):
    """
    "llmSettings": {
        "providerName": "string",
        "modelName": "string",
        "temperature": "decimal",
        "maxTokens": "integer",
        "uploadFiles": "boolean",
        "llmOutputGuardrail": "boolean",
        "inputModerationGuardrail": "boolean",
        "promptInjectionGuardrail": "boolean"
      }
    """
    provider_name: Optional[str] = None
    model_name: Optional[str] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    frequency_penalty: Optional[float] = None
    presence_penalty: Optional[float] = None
    upload_files: Optional[bool] = None
    guardrail_settings: Optional[GuardrailSettings] = None

    def to_dict(self):
        settings = {
            "providerName": self.provider_name,
            "modelName": self.model_name,
            "temperature": self.temperature,
            "maxTokens": self.max_tokens,
            "frequencyPenalty": self.frequency_penalty,
            "presence_penalty": self.presence_penalty,
            "uploadFiles": self.upload_files,
        }
        guardrail_settings = self.guardrail_settings.to_dict() if self.guardrail_settings is not None else {}
        settings.update(guardrail_settings)
        return settings

    def __str__(self):
        llm_setting = self.to_dict()
        return str(llm_setting)


class WelcomeDataFeature(BaseModel):
    """
    {
        "title": "string",
        "description": "string"
    }
    """
    title: str
    description: str

    def to_dict(self):
        return {
            "title": self.title,
            "description": self.description
        }

    def __str__(self):
        feature = self.to_dict()
        return str(feature)


class WelcomeDataExamplePrompt(BaseModel):
    """
    {
        "title": "string",
        "description": "string",
        "promptText": "string"
    }
    """
    title: str
    description: str
    prompt_text: str

    def to_dict(self):
        return {
            "title": self.title,
            "description": self.description,
            "promptText": self.prompt_text
        }

    def __str__(self):
        example_prompt = self.to_dict()
        return str(example_prompt)


class WelcomeData(BaseModel):
    """
    "title": "string",
    "description": "string",
    "features": [
        ],
        "examplesPrompt": [
        ]
      }
    """
    title: Optional[str]
    description: Optional[str]
    features: Optional[list[WelcomeDataFeature]] = []
    examples_prompt: Optional[list[WelcomeDataExamplePrompt]] = []

    def to_dict(self):
        return {
            "title": self.title,
            "description": self.description,
            "features": [feature.to_dict() for feature in self.features],
            "examplesPrompt": [example.to_dict() for example in self.examples_prompt]
        }

    def __str__(self):
        welcome_data = self.to_dict()
        return str(welcome_data)


class ChatMessage(BaseModel):
    role: str
    content: str
    function_call: Optional[Any] = None
    refusal: Optional[Any] = None
    tool_calls: Optional[Any] = None

    def to_dict(self):
        message = {
            "role": self.role,
            "content": self.content
        }
        if self.function_call:
            message["function_call"] = self.function_call

        if self.refusal:
            message["refusal"] = self.refusal

        if self.tool_calls:
            message["tool_calls"] = self.tool_calls

        return message

    def __str__(self):
        message = self.to_dict()
        return str(message)


class ChatMessageList(BaseModel):
    messages: list[ChatMessage]

    def to_list(self):
        messages = list()
        for message in self.messages:
            messages.append(message.to_dict())

        return messages


class ChatVariable(BaseModel):
    key: str
    value: str

    def to_dict(self):
        return {
            "key": self.key,
            "value": self.value
        }

    def __str__(self):
        variable = self.to_dict()
        return str(variable)


class ChatVariableList(BaseModel):
    variables: list[ChatVariable]

    def to_list(self):
        variables = list()
        for variable in self.variables:
            variables.append(variable.to_dict())

        return variables


class Assistant(BaseModel):
    """
    {
      "assistantId": "string",
      "assistantName": "string",
      "intents": [ /* full option */

      ]
    }
    """
    id: Optional[str] = None
    name: Optional[str] = None
    description: Optional[str] = None
    status: Optional[Literal[1, 2]] = 1
    priority: Optional[int] = 0
    type: Optional[str] = None
    prompt: Optional[str] = None
    intents: Optional[list[AssistantIntent]] = []
    project: Optional[Project] = None
    welcome_data: Optional[WelcomeData] = None
    llm_settings: Optional[LlmSettings] = None

    def __str__(self):
        assistant = {
            "assistantId": self.id,
            "assistantName": self.name,
            "assistantType": self.type,
            "assistantPriority": self.priority,
            "assistantStatus": self.status
        }
        if self.intents is not None and any(self.intents):
            assistant["intents"] = self.intents

        if self.description is not None:
            assistant['assistantDescription'] = self.description

        if self.prompt is not None:
            assistant['prompt'] = self.prompt

        if self.project is not None:
            assistant['project'] = self.project

        if self.welcome_data is not None:
            assistant['welcomeData'] = self.welcome_data

        if self.llm_settings is not None:
            assistant['llmSettings'] = self.llm_settings

        return str(assistant)


class TextAssistant(Assistant):
    type: Literal["text"] = "text"


class ChatAssistant(Assistant):
    type: Literal["chat"] = "chat"


class DataAnalystAssistant(Assistant):
    pass


class ChatWithDataAssistant(Assistant):
    pass


class LlmOptions(BaseModel):
    cache: Optional[bool] = False
    frequency_penalty: Optional[float] = None
    max_tokens: Optional[int] = None
    model_name: Optional[str] = None
    n: Optional[int] = None
    presence_penalty: Optional[float] = None
    provider: Optional[str] = None
    stream: Optional[bool] = False
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    type: Optional[str] = ""
    verbose: Optional[bool] = False

    def to_dict(self):
        options = {
            'cache': self.cache,
            'stream': self.stream,
            'verbose': self.verbose
        }
        if self.frequency_penalty is not None:
            options['frequencyPenalty'] = self.frequency_penalty

        if self.max_tokens is not None:
            options['maxTokens'] = self.max_tokens

        if self.model_name is not None:
            options['modelName'] = self.model_name

        if self.n is not None:
            options['n'] = self.n

        if self.presence_penalty is not None:
            options['presencePenalty'] = self.presence_penalty

        if self.provider is not None:
            options['provider'] = self.provider

        if self.temperature is not None:
            options['temperature'] = self.temperature

        if self.top_p is not None:
            options["topP"] = self.top_p

        if self.type is not None:
            options['type'] = self.type

        return options


class Search(BaseModel):
    k: int
    type: Literal["similarity", "mmr"] = "similarity"
    fetch_k: Optional[float] = None  # Valid when using mmr type
    lambda_: Optional[float] = None  # Valid when using mmr type
    prompt: str
    return_source_documents: bool
    score_threshold: float
    template: str

    def to_dict(self):
        search = {
            "k": self.k,
            'type': self.type,
            'prompt': self.prompt,
            'returnSourceDocuments': self.return_source_documents,
            'scoreThreshold': self.score_threshold,
            'template': self.template
        }
        if self.fetch_k is not None:
            search['fetchK'] = self.fetch_k

        if self.lambda_ is not None:
            search['lambda'] = self.lambda_

        return search


class RetrieverOptions(BaseModel):
    type: Literal["vectorStore", "multiQuery", "selfQuery", "hyde", "contextualCompression"]
    search_type: Optional[str] = "similarity"  # Azure AISearch specific, defaults to similarity
    step: Optional[Literal["all", "documents"]] = "all"
    prompt: Optional[str] = None  # Not needed when using vectorStore

    def to_dict(self):
        options = {
            'type': self.type,
            'step': self.step,
        }
        if self.search_type is not None:
            options['searchType'] = self.search_type

        if self.prompt is not None:
            options['prompt'] = self.prompt

        return options


class ChainOptions(BaseModel):
    type: str


class EmbeddingsOptions(BaseModel):
    dimensions: int
    model_name: str
    provider: str
    use_proxy: bool


class IngestionOptions(BaseModel):
    geai_options: Dict
    llama_parse_options: Dict
    provider: str


class SearchOptions(BaseModel):
    history_count: int
    llm: LlmOptions
    search: Search
    retriever: RetrieverOptions
    chain: Optional[ChainOptions] = None
    embeddings: Optional[EmbeddingsOptions] = None
    ingestion: Optional[IngestionOptions] = None
    options: Optional[Dict] = None
    rerank: Optional[Dict] = None
    variables: Optional[ChatVariableList] = None
    vector_store: Optional[Dict] = None

    def to_dict(self):
        options = {
            'historyCount': self.history_count,
            'llm': self.llm.to_dict() if self.llm else None,
            'search': self.search.to_dict() if self.search else None,
            'retriever': self.retriever.to_dict() if self.retriever else None
        }

        if self.chain is not None:
            options['chain'] = self.chain

        if self.embeddings is not None:
            options['embeddings'] = self.embeddings

        if self.ingestion is not None:
            options['ingestion'] = self.ingestion

        if self.options is not None:
            options['options'] = self.options

        if self.rerank is not None:
            options['rerank'] = self.rerank

        if self.variables is not None:
            options['variables'] = self.variables

        if self.vector_store is not None:
            options['vectorStore'] = self.vector_store

        return options


class ChildOptions(BaseModel):
    chunk_size: float
    chunk_overlap: float
    content_processing: Optional[Literal["", "clean"]] = ""

    def to_dict(self):
        options = {
            'chunkSize': self.chunk_size,
            'chunkOverlap': self.chunk_overlap,
        }
        if self.content_processing is not None:
            options['contentProcessing'] = self.content_processing

        return options


class ChildDocumentOptions(BaseModel):
    child_k: float
    child: ChildOptions

    def to_dict(self):
        return {
            'childK': self.child_k,
            'child': self.child.to_dict() if self.child else None
        }


class ChunkOptions(BaseModel):
    chunk_overlap: int
    chunk_size: int

    def to_dict(self):
        return {
            'chunkOverlap': self.chunk_overlap,
            'chunkSize': self.chunk_size
        }


class IndexOptions(BaseModel):
    chunks: ChunkOptions
    use_parent_document: Optional[bool] = False
    child_document: Optional[ChildDocumentOptions] = None  # Valid if use_parent_document is true

    def to_dict(self):
        options = {
            'chunks': self.chunks.to_dict(),
        }
        if self.use_parent_document is not None:
            options['useParentDocument'] = self.use_parent_document

        if self.child_document is not None:
            options['childDocument'] = self.child_document.to_dict()

        return options


class RAGAssistant(Assistant):
    template: Optional[str] = None
    search_options: Optional[SearchOptions] = None
    index_options: Optional[IndexOptions] = None
    welcome_data: Optional[WelcomeData] = None
