from abc import ABC, abstractmethod

import numpy as np
import pandas as pd
import json
import time
from IPython.core.getipython import get_ipython
from ..utils import *

class Drag(ABC):
    def  __init__(self, config):
        self.config = config

    def data(self, file_path, df_name='df'):
        self.add_cell(f"# Code generated by DRAG\nimport numpy as np\nimport pandas as pd\n\n{df_name} = pd.read_csv('{file_path}')")
        self.add_questions(question=df_name, output=str(pd.read_csv(file_path).head(3)))

    def ask(self, data: dict, question: str, context=True, question_enhancer=False):
        llm_response = self.generate_prompt(data=data, question=question, context=context, question_enhancer=question_enhancer)
        self.add_questions(question=question, output=llm_response)
        self.add_cell(f"# Code generated by DRAG\n{llm_response}")

    def add_cell(self, text, source='set_next_input', replace=False):
        shell = get_ipython()
        payload = dict(
            source=source,
            text=text,
            replace=replace,
        )
        shell.payload_manager.write_payload(payload, single=False)

    def generate_prompt(self, data, question: str, context=True, question_enhancer=False, **kwargs) -> str:
            if type(context) is bool:
                if context:
                    question_sql_list = self.get_similar_question_sql(question, **kwargs)
                    context = question_sql_list[0]
                else:
                    context = None
            if question_enhancer:
                question = self.question_enhancer(question=question, data=data, context=context)
            prompt = self.get_prompt(
                data=data,
                question=question,
                context=context,
                **kwargs,
            )
            #print(f"\n{prompt}\n")
            llm_response = self.submit_prompt(prompt, **kwargs)
            return llm_response

    def get_prompt( self, data: dict, question: str, context=None, **kwargs,):
            data_keys = [k for k in data.keys()]
            data_values = [v for v in data.values()]
            initial_prompt = f"You are an expert data scientist who helps do data analysis and write python code according to the user's request. \
It is very important that you only answer with python code and not with explanations.\n\
The code you generate will be inserted inside a Python notebook, \
in which a Pandas DataFrame called {', '.join(data_keys)} has already been declared."
            data_prompt = f"Consider the following DataFrame {', '.join(data_keys)}:\n"
            data_prompt = self.add_df_to_prompt(initial_prompt=initial_prompt, data=data_values)
            context_prompt = ""
            if context is not None and context != "":
                context_prompt = f"Consider this python code already in the notebook, use it as a basis for generating other code but it is very important not to return it as output. \
It is important to avoid duplicate code.\n'''python\n{context}\n'''"
            prompt = [
                {
                    "role": "system",
                    "content": initial_prompt
                },
                {
                    "role": "system",
                    "content": data_prompt
                },
                {
                    "role": "system",
                    "content": context_prompt
                },
                {
                    "role": "user",
                    "content": f"Only generate python code that allows: {question}"
                }
            ]
            print(f"\n\nSYSTEM:\n{initial_prompt}\n\nUSER:\nGenerates python code that enables:\n{question}")
            return prompt

    def add_df_to_prompt(self, initial_prompt: str, data) -> str:
            for d in data:
                if len(d) > 4:
                    initial_prompt += f"{str(d.head(3))}\n"
                else:
                    initial_prompt += f"{str(d)}\n"
            return initial_prompt
    
    def question_enhancer(self, question, data, context):
        data_keys = [k for k in data.keys()]
        data_values = [v for v in data.values()]
        initial_prompt = f"The user passes you a sentence that will be used as a prompt, you enhance that prompt and return it as output.\
Enhance the sentence considering that it refers to the following DataFrames {', '.join(data_keys)}\
, but it is very important not to include it in the output sentence:\n"
        initial_prompt += self.add_df_to_prompt(initial_prompt=initial_prompt, data=data_values)
        if context is not None:
            initial_prompt += f"\nAlso consider this code, rewrite the prompt so that it can be used as the basis for later generated code:\n{context}\n"
        initial_prompt  += f"\nEnhance the sentence so that the LLM receiving the prompt will be facilitated in carrying out its task.\n\
Return only the improved sentence, it is very important that you do not return anything else."
        prompt = [
                {
                    "role": "system",
                    "content": initial_prompt
                },
                {
                    "role": "user",
                    "content": f"Enhance this prompt:\n{question}"
                }
            ]
        return self.submit_prompt(prompt, temperature=0.2)
    
    def pipe(self, data: dict, question: str, context=False, question_enhancer=True):
        questions = question.split(";")
        cells = []
        llm_response = context
        for q in questions:
            llm_response = self.generate_prompt(data=data, question=q, context=llm_response, question_enhancer=question_enhancer)
            self.add_questions(question=question, output=llm_response)
            cells.append(llm_response)
        for cell in reversed(cells):
            self.add_cell(f"# Code generated by DRAG\n{cell}")

    def flow(self, data: dict, flow: str, context=False, question_enhancer=True):
        flow, mods = split_on_plus_minus(flow)
        with open(f"../src/drag/flows/{flow}.json", 'r') as f:
            flow_steps = json.load(f)
        steps_to_do = {} 
        for idx, key in enumerate(flow_steps['flow']):
            if str(idx) not in mods and key not in mods:
                steps_to_do[key] = flow_steps['flow'][key]
        cells = []
        llm_response = context
        print(steps_to_do)
        for tag, step in steps_to_do.items():
            llm_response = self.generate_prompt(data=data, question=step['Question'], context=llm_response, question_enhancer=question_enhancer)
            self.add_questions(question=step['Question'], output=llm_response)
            cells.append(llm_response)
        for cell in reversed(cells):
            self.add_cell(f"# Code generated by DRAG\n{cell}")

    def flow_steps(self, flow: str):
        with open(f"../src/drag/flows/{flow}.json", 'r') as f:
            flow_steps = json.load(f)
        print(f"{flow}\n{flow_steps['infos']['Description']}")
        for idx, key in enumerate(flow_steps['flow']):
            print(f"\n{idx} - {key}")
            for  k, v in flow_steps['flow'][key].items():
                print(f"    {k}: {v}")