import csv
from ctypes import Array
from typing import Any, Coroutine, List, Tuple
import io
import time
import re
import os

from fastapi import UploadFile
import asyncio 
import langchain
from langchain.chat_models import ChatOpenAI
from langchain.agents import create_csv_agent, load_tools, initialize_agent, AgentType, create_pandas_dataframe_agent
from langchain.tools import HumanInputRun, PythonAstREPLTool
from langchain.callbacks.tracers import ConsoleCallbackHandler
from langchain.callbacks import HumanApprovalCallbackHandler
from langchain.memory import ConversationBufferMemory, ConversationSummaryBufferMemory
from langchain import PromptTemplate
import pandas as pd
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from util.tools import SessionHumanInputRun
import util.config as config
from util.model import TemplateMappingList, TemplateMapping, TemplateMappingCode, TransformValue
import redis

r = redis.from_url(os.environ.get("REDIS_URL"))
#r = redis.from_url('redis://:password@localhost:6379')
class Processor:
    def __init__(self, session):
        self.session = session
    async def extract_csv_description(self, df: UploadFile|str, llm, memory) -> Coroutine[Any, Any, Tuple[pd.DataFrame, str]] :
        df = pd.read_csv(df)
        agent = create_pandas_dataframe_agent(llm=llm,df=df, agent_executor_kwargs={'handle_parsing_errors':True, 'memory':memory},
                                            early_stopping_method="generate", verbose=True,
                                            temperature=0,agent_type=AgentType.OPENAI_FUNCTIONS,)
        descriptions = agent.run("""Describe what is the column name of each of the column table in detail in the following format:
                            <name of column 1>: <description of column 1>\n
                            <name of column 2>: <description of column 2>""", callbacks=[ConsoleCallbackHandler()])
        return df, descriptions
    async def _human_prompt(prompt, session):
        r.publish(f'human_prompt_{session}', prompt)
    
    async def _human_input(session):
        p = r.pubsub(ignore_subscribe_messages=True)
        p.subscribe(f'human_input_{session}')
        message = None
        while True:
            message = p.get_message()
            if message and message['type']=='message':
                break
            print("waiting for human input")
            await asyncio.sleep(1)
        return message['data'].decode('utf-8')

    async def process_files(self, table_file, template_file, file_guid):
        table_string = table_file.decode('utf-8')
        template_string = template_file.decode('utf-8')
        llm = ChatOpenAI(openai_api_key=config.OPENAI_API_KEY, temperature=0, model="gpt-3.5-turbo-0613", )
        memory = ConversationSummaryBufferMemory(llm=llm,memory_key="chat_history", return_messages=True, max_token_limit=1500)
        
        table_df, table_descriptions = await self.extract_csv_description(io.StringIO(table_string), llm, memory=memory)
        r.publish(f'{self.session}_response', 'table_descriptions')
        r.publish(f'{self.session}_response', table_descriptions)

        template_df, template_descriptions = await self.extract_csv_description(io.StringIO(template_string), llm, memory=memory)
        r.publish(f'{self.session}_response', 'template_descriptions')
        r.publish(f'{self.session}_response', template_descriptions)
        dfs =[table_df, template_df]
        human_tool = SessionHumanInputRun(session=self.session)
        human_tool.description = '''
        Use this tool to take human input. 
        If the mapping is ambiguous, ask 'human' a question with options in the following format. 
        Make the human confirm the mapping by selecting the appropriate number.
            - Question: The template column <template column name> should be mapped to which one of the table columns 
            (1: <table column name 1>, 2: <table column name 2> (Recommended), 3:<table column name 3>, ...)? Select the appropriate number or specify the column name. 
        '''
        human_tool.prompt_func= Processor._human_prompt
        human_tool.input_func = Processor._human_input
        
        mappings = await self.get_mappings(llm, table_descriptions, template_descriptions, human_tool)
        codes = await self.get_template_formatting_code(llm, table_df, template_df, human_tool, mappings, memory)
        new_table_df = table_df.loc[:,[code.table_column for code in codes]]
        for code in codes:
            new_table_df[code.table_column].apply(lambda x: self.format_value(x,code=code.code))
        r.set(f"{self.session}_{file_guid}", new_table_df.to_msgpack(compress='zlib'))
        r.publish(f'{self.session}_response', f'file_guid:{file_guid}')

    
    def format_value(self, source_value, code):
        value = TransformValue(source=source_value,destination=source_value)
        try:
            exec(code, {'value':value})
        except Exception as e:
            r.publish(f'{self.session}_response',f'ERROR: \nCode: \n {code} \n Failed with error: \n{e}')
            print(e)
        return value.destination
    
    async def get_mappings(self,llm, table_descriptions, template_descriptions, human_tool):
        parser = PydanticOutputParser(pydantic_object=TemplateMappingList)
        new_parser = OutputFixingParser.from_llm(parser=parser,llm=llm)
        agent = initialize_agent(
            [human_tool],
            llm,
            agent=AgentType.OPENAI_FUNCTIONS,
            handle_parsing_errors=True,
            early_stopping_method="force",
            temperature=0.3,
            output_parser=new_parser,
        )
        descriptions = await agent.arun("""Map all the columns of the Template descriptions to columns of the table Descriptions:                          
                                    - Table Descriptions:
                                    """ + table_descriptions + """
                                    - Template Descriptions:
                                    """ + template_descriptions + """
                                    Use the table and template descriptions above to determine the mapping based on similarity, formats and distribution.
                                    If the table column names are ambiguous take human input.
                                """,callbacks=[ConsoleCallbackHandler()],)
        print(descriptions)
        
        mappings = new_parser.parse(descriptions)
        return mappings
    

    async def get_template_formatting_code(self, llm, table_df, template_df, human_tool, mappings: TemplateMappingList, memory):
        
        dfs = []
        dfs.append(table_df)
        dfs.append(template_df)
        df_locals = {}
        df_locals[f"table_df"] = table_df
        df_locals[f"template_df"] = template_df
        parser = PydanticOutputParser(pydantic_object=TemplateMappingCode)
        new_parser = OutputFixingParser.from_llm(parser=parser,llm=llm)
        
        
        codes=[]
        #The code should be in the format of a Python function taking as input a string and returning a string. 
        
        for mapping in mappings.template_mappings:
            human_tool.description = f'''
                Use this tool to get human approval. Always show the samples and code. The human can edit the code and approve it.
            '''
            table_df_samples = table_df[mapping.table_column].sample(5).to_list()
            template_df_samples = template_df[mapping.template_column].sample(5).to_list()
            agent = initialize_agent(
                        [PythonAstREPLTool(locals=df_locals)],
                        llm,
                        agent=AgentType.OPENAI_FUNCTIONS,
                        handle_parsing_errors=True,
                        early_stopping_method="force",
                        temperature=0.3,
                        output_parser=new_parser,
                        memory = memory,
                        memory_key = 'chat_history'
                    )
            #The AI can determine the format of the column values only after sampling.
            #As shown in the output below, generate the code as a Python function taking as input a string and returning a string and also include a call to the generated function.
                                
            code = agent.run(f'''Provide the code to bring the format of values in table_df column  '{mapping.table_column}'
                                to the format of values in template_df column '{mapping.template_column}' based off the values, data types and formats. 
                                Additional samples to be used to generate the code:
                                    '{mapping.table_column}' sample values: [{table_df_samples}]
                                    '{mapping.template_column}' samples values: [{template_df_samples}]
                                The input to the code will be a value object with the following attributes:
                                - source: The value of the table_df column '{mapping.table_column}'.
                                - destination: The value of the template_df column '{mapping.template_column}'.
                                Show the sample values using which the code is generated. 
                                For example, for date columns, they may be in different formats, and it is necessary to change the format from dd.mm.yyyy to mm.dd.yyyy.
                                
                                Final Answer:
                                ```
                                    ```python
                                    def format_value(source_value):
                                        <code to transform source_value into destination_value>
                                        return destination_value
                                    value.destination = format_value(value.source)
                                    ```
                                ```
                                Final Answer should contain the samples and code.
                                ''', callbacks=[ConsoleCallbackHandler(), ])
            print(code)
            human_code = await human_tool.arun(code + '\nSpecify the code with ```python``` tags.')
            regex = r"```python((.|\n|\t)*?)```"
            code = human_code if re.match(regex, human_code) else code
            
            matches = re.findall(regex, code)
            code = ''
            for match in matches:
                code = code + '\n'+ '\n'.join(match)
            codes.append(TemplateMappingCode(template_column=mapping.template_column, 
                                             table_column=mapping.table_column, 
                                             code=code))
        return codes
langchain.agents.initialize_agent langchain.memory.ConversationSummaryBufferMemory langchain.output_parsers.PydanticOutputParser langchain.tools.PythonAstREPLTool langchain.agents.create_pandas_dataframe_agent langchain.chat_models.ChatOpenAI langchain.callbacks.tracers.ConsoleCallbackHandler langchain.output_parsers.OutputFixingParser.from_llm
[((963, 990), 'os.environ.get', 'os.environ.get', (['"""REDIS_URL"""'], {}), "('REDIS_URL')\n", (977, 990), False, 'import os\n'), ((1270, 1285), 'pandas.read_csv', 'pd.read_csv', (['df'], {}), '(df)\n', (1281, 1285), True, 'import pandas as pd\n'), ((1302, 1537), 'langchain.agents.create_pandas_dataframe_agent', 'create_pandas_dataframe_agent', ([], {'llm': 'llm', 'df': 'df', 'agent_executor_kwargs': "{'handle_parsing_errors': True, 'memory': memory}", 'early_stopping_method': '"""generate"""', 'verbose': '(True)', 'temperature': '(0)', 'agent_type': 'AgentType.OPENAI_FUNCTIONS'}), "(llm=llm, df=df, agent_executor_kwargs={\n    'handle_parsing_errors': True, 'memory': memory}, early_stopping_method\n    ='generate', verbose=True, temperature=0, agent_type=AgentType.\n    OPENAI_FUNCTIONS)\n", (1331, 1537), False, 'from langchain.agents import create_csv_agent, load_tools, initialize_agent, AgentType, create_pandas_dataframe_agent\n'), ((2683, 2779), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'openai_api_key': 'config.OPENAI_API_KEY', 'temperature': '(0)', 'model': '"""gpt-3.5-turbo-0613"""'}), "(openai_api_key=config.OPENAI_API_KEY, temperature=0, model=\n    'gpt-3.5-turbo-0613')\n", (2693, 2779), False, 'from langchain.chat_models import ChatOpenAI\n'), ((2794, 2909), 'langchain.memory.ConversationSummaryBufferMemory', 'ConversationSummaryBufferMemory', ([], {'llm': 'llm', 'memory_key': '"""chat_history"""', 'return_messages': '(True)', 'max_token_limit': '(1500)'}), "(llm=llm, memory_key='chat_history',\n    return_messages=True, max_token_limit=1500)\n", (2825, 2909), False, 'from langchain.memory import ConversationBufferMemory, ConversationSummaryBufferMemory\n'), ((3498, 3540), 'util.tools.SessionHumanInputRun', 'SessionHumanInputRun', ([], {'session': 'self.session'}), '(session=self.session)\n', (3518, 3540), False, 'from util.tools import SessionHumanInputRun\n'), ((4868, 4929), 'util.model.TransformValue', 'TransformValue', ([], {'source': 'source_value', 'destination': 'source_value'}), '(source=source_value, destination=source_value)\n', (4882, 4929), False, 'from util.model import TemplateMappingList, TemplateMapping, TemplateMappingCode, TransformValue\n'), ((5288, 5345), 'langchain.output_parsers.PydanticOutputParser', 'PydanticOutputParser', ([], {'pydantic_object': 'TemplateMappingList'}), '(pydantic_object=TemplateMappingList)\n', (5308, 5345), False, 'from langchain.output_parsers import PydanticOutputParser, OutputFixingParser\n'), ((5367, 5418), 'langchain.output_parsers.OutputFixingParser.from_llm', 'OutputFixingParser.from_llm', ([], {'parser': 'parser', 'llm': 'llm'}), '(parser=parser, llm=llm)\n', (5394, 5418), False, 'from langchain.output_parsers import PydanticOutputParser, OutputFixingParser\n'), ((5434, 5614), 'langchain.agents.initialize_agent', 'initialize_agent', (['[human_tool]', 'llm'], {'agent': 'AgentType.OPENAI_FUNCTIONS', 'handle_parsing_errors': '(True)', 'early_stopping_method': '"""force"""', 'temperature': '(0.3)', 'output_parser': 'new_parser'}), "([human_tool], llm, agent=AgentType.OPENAI_FUNCTIONS,\n    handle_parsing_errors=True, early_stopping_method='force', temperature=\n    0.3, output_parser=new_parser)\n", (5450, 5614), False, 'from langchain.agents import create_csv_agent, load_tools, initialize_agent, AgentType, create_pandas_dataframe_agent\n'), ((6899, 6956), 'langchain.output_parsers.PydanticOutputParser', 'PydanticOutputParser', ([], {'pydantic_object': 'TemplateMappingCode'}), '(pydantic_object=TemplateMappingCode)\n', (6919, 6956), False, 'from langchain.output_parsers import PydanticOutputParser, OutputFixingParser\n'), ((6978, 7029), 'langchain.output_parsers.OutputFixingParser.from_llm', 'OutputFixingParser.from_llm', ([], {'parser': 'parser', 'llm': 'llm'}), '(parser=parser, llm=llm)\n', (7005, 7029), False, 'from langchain.output_parsers import PydanticOutputParser, OutputFixingParser\n'), ((10532, 10555), 're.findall', 're.findall', (['regex', 'code'], {}), '(regex, code)\n', (10542, 10555), False, 'import re\n'), ((2425, 2441), 'asyncio.sleep', 'asyncio.sleep', (['(1)'], {}), '(1)\n', (2438, 2441), False, 'import asyncio\n'), ((2988, 3013), 'io.StringIO', 'io.StringIO', (['table_string'], {}), '(table_string)\n', (2999, 3013), False, 'import io\n'), ((3250, 3278), 'io.StringIO', 'io.StringIO', (['template_string'], {}), '(template_string)\n', (3261, 3278), False, 'import io\n'), ((10459, 10486), 're.match', 're.match', (['regex', 'human_code'], {}), '(regex, human_code)\n', (10467, 10486), False, 'import re\n'), ((10690, 10801), 'util.model.TemplateMappingCode', 'TemplateMappingCode', ([], {'template_column': 'mapping.template_column', 'table_column': 'mapping.table_column', 'code': 'code'}), '(template_column=mapping.template_column, table_column=\n    mapping.table_column, code=code)\n', (10709, 10801), False, 'from util.model import TemplateMappingList, TemplateMapping, TemplateMappingCode, TransformValue\n'), ((1905, 1929), 'langchain.callbacks.tracers.ConsoleCallbackHandler', 'ConsoleCallbackHandler', ([], {}), '()\n', (1927, 1929), False, 'from langchain.callbacks.tracers import ConsoleCallbackHandler\n'), ((7662, 7697), 'langchain.tools.PythonAstREPLTool', 'PythonAstREPLTool', ([], {'locals': 'df_locals'}), '(locals=df_locals)\n', (7679, 7697), False, 'from langchain.tools import HumanInputRun, PythonAstREPLTool\n'), ((6408, 6432), 'langchain.callbacks.tracers.ConsoleCallbackHandler', 'ConsoleCallbackHandler', ([], {}), '()\n', (6430, 6432), False, 'from langchain.callbacks.tracers import ConsoleCallbackHandler\n'), ((10224, 10248), 'langchain.callbacks.tracers.ConsoleCallbackHandler', 'ConsoleCallbackHandler', ([], {}), '()\n', (10246, 10248), False, 'from langchain.callbacks.tracers import ConsoleCallbackHandler\n')]