import json
import importlib
from langchain import PromptTemplate

class AspectFactory:
    def __init__(self, llm=None, ddb_util=None, google_util=None, context=None, aspect=None):

        DDBUtilWrapper = importlib.import_module('autopop.ddb_util_wrapper').DDBUtilWrapper
        GoogleUtilWrapper = importlib.import_module('autopop.google_util_wrapper').GoogleUtilWrapper

        self.llm = llm
        self.aspect = aspect
        self.context = context
        self.gu = google_util
        
        # load data from DDB
        try:
            self.template = ddb_util.get_prompt_data(context=self.context, aspect=aspect)        
            self.prompt = self.template.get('prompt', 'No prompt provided.')
        except Exception as e:
            print(e)
        
    def get_prompt(self, use_prompt=False):
        if use_prompt: return self.prompt

        template = """{guidelines}

Setting: {setting}

Persona: {persona}

Tone: {tone}

Task: {task}

Form: {form}

Examples: {examples}

Do not hallucinate - use only the context provided.

Answer: """

        prompt_template = PromptTemplate(
            input_variables=["guidelines", "setting", "persona", "task", "tone", "form", "examples"],
            template=template
        )

        formatted_prompt = prompt_template.format(
            guidelines=self.template.get('guidelines', 'No guidelines provided.'),
            setting=self.template.get('setting', 'No setting provided.'),
            persona=self.template.get('persona', 'No persona provided.'),
            task=self.template.get('task', 'No task provided.'),
            tone=self.template.get('tone', 'No tone provided.'),
            form=self.template.get('form', 'No form provided.'),
            examples='\n'.join(self.template.get('examples', ['No examples provided.']))
        )
        self.prompt = formatted_prompt
        return formatted_prompt

    def query(self, use_prompt=False, max_retries=3):
        retries = 0
        while retries < max_retries:
            prompt = self.get_prompt(use_prompt=use_prompt)
            answer = self.llm.query(prompt).answer
            if answer.lower() != 'i cannot answer.': return answer
            retries += 1
        return answer

    def post_doc(self, destination=None, use_prompt=False):
        fields = [{'table':self.context, 'row': self.aspect}]
        self.gu.post_doc_data(fields=fields, dest=destination, data=self.query(use_prompt=use_prompt)) 

    def post_sheet_query(self, sheet=None, column='D', q=query):
        fields = [{'table':self.context, 'row': self.aspect}]
        self.gu.post_sheet_data(fields=fields, sheet=sheet, column=column, data=query) 

    def post_sheet(self, sheet=None, column='D', use_prompt=False):
        fields = [{'table':self.context, 'row': self.aspect}]
        self.gu.post_sheet_data(fields=fields, sheet=sheet, column=column, data=self.query(use_prompt=use_prompt)) 

    def __str__(self):
        return json.dumps(self.template)
        
####################################################################

# fix URL
class DDBPromptWrapper(AspectFactory):
    def __init__(self, ddb_util=None, google_util=None, llm=None, context=None, aspect=None):
        super().__init__(ddb_util=ddb_util, google_util=google_util, llm=llm, context=context, aspect=aspect)

####################################################################

# Function to initialize the dynamic classes
def dynamic_init(self, context, aspect):
    super(self.__class__, self).__init__(context=context, aspect=aspect)   
    
# Function to get instances
def get_instances(aspects):
    # Create classes dynamically
    for aspect in aspects:
        for row in aspect['row']:
            class_name = aspect['table'] + row
            globals()[class_name] = type(class_name, (DDBP,), {"__init__": lambda self, context=aspect['table'], aspect=row: dynamic_init(self, context, aspect)})

    # Create instances dynamically
    return [globals()[aspect['table'] + row]() for aspect in aspects for row in aspect['row']]  

