from langchain_groq import ChatGroq
import os
from dotenv import load_dotenv
import pandas as pd
import json

class FewShotPosts:
    def __init__(self, file_path="processed_posts.json"):
        self.df = None
        self.unique_tags = None
        self.load_posts(file_path)

    def load_posts(self, file_path):
        with open(file_path, encoding="utf-8") as f:
            posts = json.load(f)
            self.df = pd.json_normalize(posts)
            self.df['length'] = self.df['line_count'].apply(self.categorize_length)
            all_tags = self.df['tags'].apply(lambda x: x).sum()
            self.unique_tags = list(set(all_tags))

    def get_filtered_posts(self, length, language, tag):
        df_filtered = self.df[
            (self.df['tags'].apply(lambda tags: tag in tags)) &
            (self.df['language'] == language) &
            (self.df['length'] == length)
        ]
        return df_filtered.to_dict(orient='records')

    def categorize_length(self, line_count):
        if line_count < 5:
            return "Short"
        elif 5 <= line_count <= 10:
            return "Medium"
        else:
            return "Long"

    def get_tags(self):
        return self.unique_tags

load_dotenv()
few_shot = FewShotPosts()

def get_length_str(length):
    if length == "Short":
        return "1 to 5 lines"
    elif length == "Medium":
        return "6 to 10 lines"
    elif length == "Long":
        return "11 to 15 lines"
    elif length == "Extra Long":
        return "16 to 20 lines"
    else:
        return "6 to 10 lines"

def generate_post(length, language, tag, model_name="llama-3.3-70b-versatile", custom_context=""):
    prompt = get_prompt(length, language, tag, custom_context)
    llm_instance = ChatGroq(groq_api_key=os.getenv("GROQ_API_KEY"), model_name=model_name)
    response = llm_instance.invoke(prompt)
    return response.content

def get_prompt(length, language, tag, custom_context=""):
    length_str = get_length_str(length)

    prompt = f"""
    Generate a LinkedIn post using the below information. No preamble.
    1) Topic: {tag}
    2) Length: {length_str}
    3) Language: {language}
    If Language is Hinglish then it means it is a mix of Hindi and English. 
    The script for the generated post should always be English.
    """

    # Add custom context if provided
    if custom_context.strip():
        prompt += f"\nUse the following additional knowledge or context about the topic:\n{custom_context}\n"

    # Optionally include few-shot examples
    examples = few_shot.get_filtered_posts(length, language, tag)
    if len(examples) > 0:
        prompt += "\n4) Use the writing style as per the following examples:"
        for i, post in enumerate(examples):
            post_text = post['text']
            prompt += f"\n\nExample {i+1}:\n{post_text}"
            if i == 1:  # limit to 2 examples
                break

    return prompt

if __name__ == "__main__":
    print(generate_post("Medium", "English", "Mental Health", custom_context="Remember to mention the importance of daily meditation."))