import re
import openai
import inflect
import pandas as pd
from typing import Dict
from datasets import load_dataset
from IPython.display import display, HTML
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.vectorstores.utils import DistanceStrategy

import os

OPENAI_KEY = ''

openai.api_key = OPENAI_KEY
os.environ['OPENAI_API_KEY'] = OPENAI_KEY


# Constants
FS_COLUMNS = ['asin', 'category', 'title', 'tech_process', 'labels']
MAX_TOKENS = 700
USER_TXT = 'Write feature-bullets for an Amazon product page. ' \
           'Title: {title}. Technical details: {tech_data}.\n\n### Feature-bullets:'
           
# Load few-shot dataset
FS_DATASET = load_dataset('iarbel/amazon-product-data-filter', split='validation')

# Prepare Pandas DFs with the relevant columns
FS_DS = FS_DATASET.to_pandas()[FS_COLUMNS]

# Load vector store
DB = FAISS.load_local('data/vector_stores/amazon-product-embedding', OpenAIEmbeddings(),
                      distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT)


class Conversation:
    """
    A class to construct conversations with the ChatAPI
    """
    def __init__(self):
        self.messages = [{'role': 'system',
                          'content': 'You are a helpful assistant. Your task is to write feature-bullets for an Amazon product page.'}]

    def add_message(self, role: str, content: str) -> None:
        # Validate inputs
        role = role.lower()
        last_role = self.messages[-1]['role']
        if role not in ['user', 'assistant']:
            raise ValueError('Roles can be "user" or "assistant" only')
        if role == 'user' and last_role not in ['system', 'assistant']:
            raise ValueError('"user" message can only follow "assistant" message')
        elif role == 'assistant' and last_role != 'user':
            raise ValueError('"assistant" message can only follow "user" message')
        
        message = {"role": role, "content": content}
        self.messages.append(message)
            
    def display_conversation(self) -> None:
        SEP = '\n'
        for message in self.messages:
            if message['role'] == 'system':
                display(HTML(f'<b>{message["content"]}</b>'))
            elif message['role'] == 'user':
                msg_align = message["content"].replace("Title:", "<br><b>Title:</b>")\
                    .replace("Technical details:", "<br><b>Technical details:</b>").replace("### Feature-bullets:", "<br><b>Feature-bullets:</b>")
                display(HTML(f'<p style="background-color:White; color:Black; padding:5px;">{msg_align}</p>'))
            else:
                msg_align = message["content"].lstrip(SEP).replace(SEP, "<br><br>")
                display(HTML(f'<p style="background-color:LightGray; color:Black; padding:5px;">{msg_align}</p>'))

def api_call(messages: Dict[str, str], temperature: float = 0.7, top_p: int = 1, n_responses: int = 1) -> dict:
    """
    A function to call the ChatAPI. Taken in a conversation, and the optional params temperature (controls randomness) and n_responses
    """
    params = {'model': 'gpt-3.5-turbo', 'messages': messages, 'temperature': temperature, 'max_tokens': MAX_TOKENS, 'n': n_responses, 'top_p': top_p}
    response = openai.ChatCompletion.create(**params)

    text = [response['choices'][i]['message']['content'] for i in range(n_responses)]
    out = {'object': 'chat', 'usage': response['usage']._previous, 'text': text}
    return out


class FewShotData:
    def __init__(self, few_shot_df: pd.DataFrame, vector_db: FAISS):
        self.few_shot_df = few_shot_df
        self.vector_db = vector_db
        
    def extract_few_shot_data(self, target_title: str, k_shot: int = 2, **db_kwargs) -> pd.DataFrame:
         # Find relevant products
        target_title_vector = OpenAIEmbeddings().embed_query(target_title)
        similarity_list_mmr = self.vector_db.max_marginal_relevance_search_with_score_by_vector(target_title_vector, k=k_shot, **db_kwargs)
        few_shot_titles = [i[0].page_content for i in similarity_list_mmr]
        
        # Extract relevant data
        few_shot_data = self.few_shot_df[self.few_shot_df['title'].isin(few_shot_titles)][['title', 'tech_process', 'labels']]
        return few_shot_data

    def construct_few_shot_conversation(self, target_title: str, target_tech_data: str, few_shot_data: pd.DataFrame) -> Conversation:
        # Structure the few-shott data
        fs_titles = few_shot_data['title'].to_list()
        fs_tech_data = few_shot_data['tech_process'].to_list()
        fs_labels = few_shot_data['labels'].to_list()
    
        # Init a conversation, populate with few-shot data
        conv = Conversation()
        for title, tech_data, lables in zip(fs_titles, fs_tech_data, fs_labels):
            conv.add_message('user', USER_TXT.format(title=title, tech_data=tech_data))
            conv.add_message('assistant',lables)
            
        # Add the final user prompt
        conv.add_message('user', USER_TXT.format(title=target_title, tech_data=target_tech_data))
        return conv
    
    
def return_is_are(text: str) -> str:
    engine = inflect.engine()
    res = 'is' if not engine.singular_noun(text) else 'are'
    return res

def format_tech_as_str(tech_data):
    tech_format = [f'{k} {return_is_are(k)} {v}' for k, v in tech_data.to_numpy() if k and v]
    tech_str = '. '.join(tech_format)
    return tech_str


def generate_data(title: str, tech_process: str, few_shot_df: pd.DataFrame, vector_db: FAISS) -> str:
    fs_example = FewShotData(few_shot_df=few_shot_df, vector_db=vector_db)
    fs_data = fs_example.extract_few_shot_data(target_title=title, k_shot=2)
    
    fs_conv = fs_example.construct_few_shot_conversation(target_title=title, 
                                                         target_tech_data=tech_process, 
                                                         few_shot_data=fs_data)
    
    api_res = api_call(fs_conv.messages, temperature=0.7)
    feature_bullets = "## Feature-Bullets\n" + api_res['text'][0]
    return feature_bullets


def check_url_structure(url: str) -> bool:
    pattern = r"https://www.amazon.com(/.+)?/dp/[a-zA-Z0-9]{10}/?$"
    return bool(re.match(pattern, url))