from typing import Any, Dict
from openai import OpenAI
from string import Template 
import logging
import json


class LLMHandler:
    def __init__(self, api_key, model_name, default_temperature, default_max_tokens, default_system_message=None):

        self.client = OpenAI(api_key=api_key)
        self.model_name = model_name
        self.default_temperature = default_temperature
        self.default_max_tokens = default_max_tokens
        self.default_system_message = default_system_message
  
    def generate(
            self, 
            prompt, 
            model_name, 
            temperature=None, 
            max_tokens=None, 
            system_message=None):
        
        # optional parameters
        model_name = model_name or self.model_name
        temperature = temperature or self.default_temperature
        max_tokens = max_tokens or self.default_max_tokens
        system_message = system_message or self.default_system_message

        # prepare messages
        messages = [{"role": "user", "content": prompt}]
        if system_message:
            messages.insert(0, {"role": "system", "content": system_message}) 

        # get response
        logging.info(f"Generating text with model {model_name}, temperature {temperature}, and max tokens {max_tokens}...")
        response = self.client.chat.completions.create(
            model=model_name,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            seed=1234,
            response_format={ "type": "json_object" } 
        )
        return response.choices[0].message.content



def format_personalization(personalization):
    return '\n'.join([f"- {key}: {value}" for key, value in personalization.items()])


def build_prompt(
        context: Dict[str, Any],
        few_shot: str = None,
        ) -> str:
        """
        Create a detailed prompt incorporating all personalization factors
        
        :param section_type: Type of section to generate
        :param context: Context data including items and customer info
        :param additional_textual_context: additional context for personalization
        :return: Formatted prompt string
        """

        if few_shot is not None and len(few_shot) > 0:
            few_shot = f"Few-shot examples:\n{few_shot}\n"
        else:
            few_shot = ""

        # Create a detailed user prompt with all context and personalization
        user_prompt = f"""
        Generate personalized newsletter content based on the following information:

        Customer Purchase Context:
        - Previously bought items: {context.pop('bought_items', [])}
        - Recommended items: {context.pop('recommended_items', [])}
        
        Personalization Parameters:
        {format_personalization(context)}

        Required Sections:
        1. greeting: A warm personalized welcome that considers all personalization parameters. Start with "Hello [customer name]," 
        2. intro: Acknowledge previous purchases and build connection warmly. Mention the previous purchases. Only acknoweldge previous purchases that are actually present in the context. 
        3. recommendations: Present recommended items naturally, explaining why they match the customer's preferences and previous purchases.
        4. closing: Wrap up maintaining the established tone. 

        General requirements:
        - Write in first person. 
        - Don't be too formal, maintain a friendly, warm and engaging tone, avoid stuff like "I hope you are doing well".
        - The reader should feel like they have a special dedicated fashion assistant, who is also a friend.
        - When mentioning items, use the item name and avoid mentioning colors that are not present in the image, or using adjectives like "vibrant".
        - The goal should be to make the reader feel special and excited about the items they are being recommended.

        {few_shot}
        
        Please generate all sections in a single, coherent response that maintains consistency in tone and style throughout. Begin each section with the section name in all caps.
        """
        return user_prompt


def initialize_newsletter(newsletter_meta_info, transactions, recommendations):
    
    # get the paths to the template files
    newsletter_example_path = newsletter_meta_info['newsletter_example_path']
    
    # load the template from file
    with open(newsletter_example_path, "r") as f:
        newsletter_text = Template(f.read())
        newsletter_text = newsletter_text.safe_substitute(
            brand_logo=newsletter_meta_info.get("brand_logo", ""),
            brand_name=newsletter_meta_info.get("brand_name", ""),
            )
    
    # iterate over the transactions and replace the placeholders with the actual content
    for i, transaction in enumerate(transactions):
        newsletter_text = newsletter_text.replace("${transaction_url}", transaction['product_image_url'], 1)
        newsletter_text = newsletter_text.replace("${transaction_name}", transaction['product_name'], 1)

    # iterate over the recommendations and replace the placeholders with the actual content
    for i, recommendation in enumerate(recommendations):
        newsletter_text = newsletter_text.replace("${recommendation_url}", recommendation['product_image_url'], 1)
        newsletter_text = newsletter_text.replace("${recommendation_name}", recommendation['product_name'], 1)
    return newsletter_text


def integrate_personalized_text(newsletter_text, customer_info, textual_sections):
    # get dctionary from the textual sections json
    textual_sections = json.loads(textual_sections)

    logging.info(textual_sections)

    newsletter_text = Template(newsletter_text)
    newsletter_text = newsletter_text.safe_substitute(
        **textual_sections
    )

    # make sure the greeting line is h1, so replace Hello ${customer_name} with <h1>Hello ${customer_name}</h1>
    customer_name = customer_info.get("customer name")
    newsletter_text = newsletter_text.replace(f"Hello {customer_name},", f'<h1 style="color:black;">Hello {customer_name},</h1>')

    return newsletter_text


def build_context(recommendations, transactions, additional_context, customer_info):
    
    # clean recommendations and transactions 
    for rec in recommendations:
        rec.pop('product_image_url', None)
        rec.pop('article_id', None)
        rec.pop('sales_channel_id', None)
        rec.pop('transaction_date', None)
    for tr in transactions:
        tr.pop('product_image_url', None)
        tr.pop('article_id', None)
        tr.pop('sales_channel_id', None)
        tr.pop('transaction_date', None)
        tr.pop('price', None)    

    context = {
        'recommended_items': recommendations,
        'bought_items': transactions,
        'customer_name': customer_info.get('customer name', 'Unknown'),
        'customer_age': customer_info.get('customer age', 'Unknown'),
        'last_visit': customer_info.get('last_visit', 'Unknown'),
        # 'preferred_categories': customer_info.get('preferred_categories', 'Unknown'),
        # 'shopping_frequency': customer_info.get('shopping_frequency', 'Unknown'),
        # 'style_preferences': customer_info.get('style_preferences', 'Unknown'),
        # 'size_preferences': customer_info.get('size_preferences', 'Unknown'),
        'preferences': additional_context if len(additional_context) > 1 else 'No additional preferences',
    }

    return context