import langcodes
from flask import current_app
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.prompts import ChatPromptTemplate
import ast
from typing import List

from common.models.document import EmbeddingSmallOpenAI


class CitedAnswer(BaseModel):
    """Default docstring - to be replaced with actual prompt"""

    answer: str = Field(
        ...,
        description="The answer to the user question, based on the given sources",
    )
    citations: List[int] = Field(
        ...,
        description="The integer IDs of the SPECIFIC sources that were used to generate the answer"
    )


def set_language_prompt_template(cls, language_prompt):
    cls.__doc__ = language_prompt


def select_model_variables(tenant):
    embedding_provider = tenant.embedding_model.rsplit('.', 1)[0]
    embedding_model = tenant.embedding_model.rsplit('.', 1)[1]

    llm_provider = tenant.llm_model.rsplit('.', 1)[0]
    llm_model = tenant.llm_model.rsplit('.', 1)[1]

    # Set model variables
    model_variables = {}
    if tenant.es_k:
        model_variables['k'] = tenant.es_k
    else:
        model_variables['k'] = 5

    if tenant.es_similarity_threshold:
        model_variables['similarity_threshold'] = tenant.es_similarity_threshold
    else:
        model_variables['similarity_threshold'] = 0.7

    if tenant.chat_RAG_temperature:
        model_variables['RAG_temperature'] = tenant.chat_RAG_temperature
    else:
        model_variables['RAG_temperature'] = 0.3

    if tenant.chat_no_RAG_temperature:
        model_variables['no_RAG_temperature'] = tenant.chat_no_RAG_temperature
    else:
        model_variables['no_RAG_temperature'] = 0.5

    # Set Tuning variables
    if tenant.embed_tuning:
        model_variables['embed_tuning'] = tenant.embed_tuning
    else:
        model_variables['embed_tuning'] = False

    if tenant.rag_tuning:
        model_variables['rag_tuning'] = tenant.rag_tuning
    else:
        model_variables['rag_tuning'] = False

    # Set HTML Chunking Variables
    model_variables['html_tags'] = tenant.html_tags
    model_variables['html_end_tags'] = tenant.html_end_tags
    model_variables['html_included_elements'] = tenant.html_included_elements
    model_variables['html_excluded_elements'] = tenant.html_excluded_elements

    # Set Embedding variables
    match embedding_provider:
        case 'openai':
            match embedding_model:
                case 'text-embedding-3-small':
                    api_key = current_app.config.get('OPENAI_API_KEY')
                    model_variables['embedding_model'] = OpenAIEmbeddings(api_key=api_key,
                                                                          model='text-embedding-3-small')
                    model_variables['embedding_db_model'] = EmbeddingSmallOpenAI
                    model_variables['min_chunk_size'] = current_app.config.get('OAI_TE3S_MIN_CHUNK_SIZE')
                    model_variables['max_chunk_size'] = current_app.config.get('OAI_TE3S_MAX_CHUNK_SIZE')
                case _:
                    raise Exception(f'Error setting model variables for tenant {tenant.id} '
                                    f'error: Invalid embedding model')
        case _:
            raise Exception(f'Error setting model variables for tenant {tenant.id} '
                            f'error: Invalid embedding provider')

    # Set Chat model variables
    match llm_provider:
        case 'openai':
            api_key = current_app.config.get('OPENAI_API_KEY')
            model_variables['llm'] = ChatOpenAI(api_key=api_key,
                                                model=llm_model,
                                                temperature=model_variables['RAG_temperature'])
            tool_calling_supported = False
            match llm_model:
                case 'gpt-4-turbo' | 'gpt-4o':
                    summary_template = current_app.config.get('GPT4_SUMMARY_TEMPLATE')
                    rag_template = current_app.config.get('GPT4_RAG_TEMPLATE')
                    history_template = current_app.config.get('GPT4_HISTORY_TEMPLATE')
                    tool_calling_supported = True
                case 'gpt-3-5-turbo':
                    summary_template = current_app.config.get('GPT3_5_SUMMARY_TEMPLATE')
                    rag_template = current_app.config.get('GPT3_5_RAG_TEMPLATE')
                    history_template = current_app.config.get('GPT3_5_HISTORY_TEMPLATE')
                case _:
                    raise Exception(f'Error setting model variables for tenant {tenant.id} '
                                    f'error: Invalid chat model')
            model_variables['summary_template'] = summary_template
            model_variables['rag_template'] = rag_template
            model_variables['history_template'] = history_template
            if tool_calling_supported:
                model_variables['cited_answer_cls'] = CitedAnswer
        case _:
            raise Exception(f'Error setting model variables for tenant {tenant.id} '
                            f'error: Invalid chat provider')

    return model_variables


def create_language_template(template, language):
    try:
        full_language = langcodes.Language.make(language=language)
        language_template = template.replace('{language}', full_language.display_name())
    except ValueError:
        language_template = template.replace('{language}', language)

    return language_template