import os

from omegaconf import OmegaConf

from dotenv import load_dotenv
load_dotenv(override=True)

from pydantic import Field, BaseModel
from vectara_agent.agent import Agent
from vectara_agent.tools import ToolsFactory, VectaraToolFactory
from vectara_agent.tools_catalog import rephrase_text

teaching_styles = ['Inquiry-based', 'Socratic', 'traditional']
languages = {'English': 'en', 'Spanish': 'es', 'French': 'fr', 'German': 'de', 'Arabic': 'ar', 'Chinese': 'zh-cn', 
             'Hebrew': 'he', 'Hindi': 'hi', 'Italian': 'it', 'Japanese': 'ja', 'Korean': 'ko', 'Portuguese': 'pt'}
initial_prompt = "How can I help you today?"

def create_assistant_tools(cfg):        

    def adjust_response_to_student(
            text: str = Field(description='the text to adjust. may include citations in markdown format.'), 
            age: int = Field(description='the age of the student. An integer'), 
            style: str = Field(description='teaching style'), 
            language: str = Field(description='the language')

        ) -> str:
        """
        Rephrase the provided text to match the student's age, desired teaching style and language.
        """
        instructions = f'''
        The following is some text, which may include citations in markdown format.
        Adjust the response to match the student's age of {age}, the {style} teaching style.
        Where possible, maintain the citations in context of the response (instead of listing them at the end). 
        When showing citations, use markdown in the following format: [[N](URL)], where N is the citation sequential number.
        You can renumber citations if needed.
        For example, in the inquiry-based teaching style, choose to ask questions that encourage the student to think critically instead of repsonding directly with the answer.
        Or in the socratic teaching style, choose to ask questions that lead the student to the answer.
        Always respond in the {language} language.''' \
            .replace("{style}", cfg.style) \
            .replace("{language}", cfg.language) \
            .replace("{student_age}", str(cfg.student_age))

        return rephrase_text(text, instructions)


    class JusticeHarvardArgs(BaseModel):
        query: str = Field(..., description="The user query.")

    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key, 
                                     vectara_customer_id=cfg.customer_id, 
                                     vectara_corpus_id=cfg.corpus_id)
    tools_factory = ToolsFactory()
    query_tool = vec_factory.create_rag_tool(
        tool_name = "justice_harvard_query",
        tool_description = """
        Answer questions about the justice, morality, politics and related topics,
        based on transcripts of recordings from the Justice Harvard class that includes a lot of content on these topics.
        When using the tool it's best to ask simple short questions. You can break complex questions into sub-queries.
        """,
        tool_args_schema = JusticeHarvardArgs,
        reranker = "multilingual_reranker_v1", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 10,
        vectara_summarizer = 'vectara-summary-ext-24-05-med-omni',
        include_citations = True,
    )

    return (
            [tools_factory.create_tool(tool) for tool in
                [
                    adjust_response_to_student,
                ]
            ] +
            tools_factory.standard_tools() + 
            tools_factory.guardrail_tools() +
            [query_tool]
    )

def initialize_agent(_cfg, update_func=None):
    bot_instructions = f"""
    - You are a helpful teacher assistant, with expertise in education in various teaching styles.
    - Obtain information using tools to answer the user's query.
    - If the tool cannot provide information relevant to the user's query, tell the user that you are unable to provide an answer.
    - If the tool can provide relevant information, use the adjust_response_to_student tool 
      to rephrase the text (including citations if any) to ensure it fits the student's age of {_cfg.student_age}, 
      the {_cfg.style} teaching style and the {_cfg.language} language.
    - When including links to the Justice Harvard videos, the url should include the start time (from the metadata).
      For example: https://youtube.com/watch?v=0O2Rq4HJBxw&t=1234 represents a link to a youtube video starting at 1234 seconds. 
    - When showing citations, renumber them if needed and use markdown in the following format: [[N](URL)], where N is the citation sequential number.
    - Response in a concise and clear manner, and provide the most relevant information to the student.
    - Never discuss politics, and always respond politely.
    """

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="justice, morality, politics, and philosophy",
        custom_instructions=bot_instructions,
        update_func=update_func
    )
    return agent


def get_agent_config() -> OmegaConf:
    cfg = OmegaConf.create({
        'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
        'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
        'api_key': str(os.environ['VECTARA_API_KEY']),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_name': "Justice-Harvard",
        'demo_welcome': "Welcome to the Justice Harvard e-learning assistant demo.",
        'demo_description': "Ask questions about the Justice Harvard class and get answers from an expert teaching assistant.",
        'style': teaching_styles[0],
        'language': 'English',
        'student_age': 18
    })
    return cfg