from .dataiku_client import DataikuClient
from datetime import datetime
import os
import re
import dataiku
from dataiku.langchain.dku_llm import DKULLM
from .config import FFConfig
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from backend.models import LlmHistory
from typing import Any, List, Dict
import re
from email.message import EmailMessage
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from .setup_logger import logger
from langchain.prompts import PromptTemplate


class FFGPT:
    def __init__(self):
        self.dataiku_client = DataikuClient()
        os.environ["DKU_CURRENT_PROJECT_KEY"] = self.dataiku_client.project_key
        self.llm_id = self.dataiku_client.variables['llm_id']
        self.max_num_token = self.dataiku_client.variables['max_num_tokens']
        self.report_size = 2000
        memory_max_token_limit = self.max_num_token - self.report_size
        self.setup_llm()
        self.reports_folder = dataiku.Folder(
            FFConfig.REPORTS_FOLDER,
            project_key=self.dataiku_client.project_key,
        )
        self.saved_reports_folder = dataiku.Folder(
            FFConfig.SAVED_REPORTS_FOLDER,
            project_key=self.dataiku_client.project_key,
        )
        self.conclusion = dataiku.Dataset(
            name=FFConfig.CONCLUSION_DATASET,
            project_key=self.dataiku_client.project_key,
        ).get_dataframe()
        self.memory = ConversationBufferMemory(
            max_token_limit=memory_max_token_limit, return_messages=True
        )

    def update_memory(self, chat_history: List[LlmHistory]) -> None:
        self.memory.clear()
        for item in chat_history:
            inputs = {"input": item["query"]}
            outputs = {"output": item["email"]}

            logger.info(f"The memory inputs {inputs}")
            logger.info(f"The memory outputs {outputs}")

            self.memory.save_context(inputs, outputs)
    
    @property
    def business_units(self):
        return [{'business_unit': x} for x in list(self.conclusion['category'].unique())]

    def setup_llm(self):
        #llm_id = "azureopenai:bs-azure-openai:doc-intell"
        self.llm = DKULLM(llm_id=self.llm_id, max_tokens=self.report_size)

    def generate_email(self, business_unit):
        with self.reports_folder.get_download_stream(business_unit + '_email.html') as f:
            email = f.read().decode()
        return email
    
    def send_message(
        self,
        user_input,
        chat_history
    ):
        self.update_memory(chat_history=chat_history)
        prompt_template =  PromptTemplate.from_template(
            """The AI is an assistant editing an HTML report following the Human's instructions. 
            The AI only returns the HTML report according to the changes required. 
            If the AI does not understand the instruction, it returns the HTML report unchanged.
            \n\nCurrent conversation:\n{history}\nHuman: {input}\nAI:"""
        )
        qa = ConversationChain(llm=self.llm, memory=self.memory, verbose=True, prompt=prompt_template)
        query = """
        Edit the report following the instructions below and return it. Always return the edited report but nothing else.
        
        {user_input}

        Report:
        """.format(user_input=user_input)
        # print(self.num_tokens_from_messages())
        answer = qa.predict(input=query)
        logger.info(answer)
        pattern = r'(?s)<html>.*?</html>'
        html_content = re.search(pattern, answer)
        if html_content is not None:
            return html_content.group(0), 1
        else:
            return chat_history[-1]['email'], 0

    def save_report(self, email_content, business_unit):
        # write email content in folder
        with self.saved_reports_folder.get_writer(business_unit + '_email_' + (datetime.now()).strftime("%Y-%m-%d_%H:%M:%S") + '.html') as f:
            f.write(email_content.encode())
        return 1

    def get_saved_folder_url(self):
        instance_url = self.dataiku_client.client.get_general_settings().get_raw().get('studioExternalUrl')
        project_key = self.dataiku_client.project_key
        saved_folder_url = instance_url + '/projects/' + project_key + '/managedfolder/nwzNmkjP/view/'
        return saved_folder_url
