import os
import faiss
from sentence_transformers import SentenceTransformer
from typing import List, Dict
from openai import OpenAI
import numpy as np
import pandas as pd
import requests
import base64
from huggingface_hub import login
from helpercodes.image_functions import _encode_image_to_base64

from dotenv import load_dotenv
import os

load_dotenv()

login(token=os.getenv("HF_TOKEN"))

class ContextRetrieverWithFullText:

    def __init__(self):
        self.sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')
        self.keywords = []
        self.image_urls = []
        self.keyword_to_url_map = {}

        self.embedding_dim = 384
        self.index = faiss.IndexFlatIP(self.embedding_dim)

    def add_keywords(self, keywords: List[str], image_urls: List[str]):

        all_embeddings = []
        start_idx = len(self.keywords)

        self.keywords.extend(keywords)
        self.image_urls.extend(image_urls)

        for idx, keyword in enumerate(keywords, start=start_idx):
            embedding = self.sentence_transformer.encode(keyword)
            all_embeddings.append(embedding)
            self.keyword_to_url_map[idx] = image_urls[idx]

        if all_embeddings:
            embeddings_array = np.array(all_embeddings).astype('float32')
            self.index.add(embeddings_array)

    def find_relevant_keywords(self, query: str, n_results: int = 2) -> List[Dict]:

        query_embedding = self.sentence_transformer.encode(query)
        distances, indices = self.index.search(
            np.array([query_embedding]).astype('float32'),
            k=n_results
        )

        results = []
        for distance, idx in zip(distances[0], indices[0]):
            if idx < len(self.keywords):
                keyword = self.keywords[idx]
                image_url = self.keyword_to_url_map[idx]
                results.append({
                    'keyword': keyword,
                    'image_url': image_url,
                    'relevance_score': float(distance)
                })

        return results


def get_dataset_dict():

    script_dir = os.path.dirname(os.path.abspath(__file__))
    csv_path = os.path.join(script_dir, "newData.csv")
    df = pd.read_csv(csv_path)
    data = df[['Image Url', 'Keywords', 'styles', 'Title', 'Body', 'Description']].copy()

    unique_keywords_by_source = {}
    for source in data['styles'].unique():
        source_key = source.lower().replace(' ', '_').replace('/', '_')
        data.loc[data['styles'] == source, 'styles'] = source_key

        source_data = data[data['styles'] == source_key]
        all_keywords = source_data['Keywords'].dropna().str.lower().str.split(', ')
        flattened_keywords = [kw for sublist in all_keywords for kw in sublist]
        unique_keywords = list(set([kw for kw in flattened_keywords if source_key not in kw]))
        unique_keywords_by_source[source_key] = unique_keywords

    return data, unique_keywords_by_source

def get_links(df: pd.DataFrame, input_keywords: str):

    retriever = ContextRetrieverWithFullText()
    keywords = df["Keywords"].tolist()
    image_urls = df["Image Url"].tolist()

    retriever.add_keywords(keywords, image_urls)

    results = retriever.find_relevant_keywords(input_keywords, n_results=2)
    return results


def generate_gpt_prompt_our_data(results):

    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    response1 = requests.get(results[0]['image_url'])
    base64_image1 = base64.b64encode(response1.content).decode('utf-8')

    response2 = requests.get(results[1]['image_url'])
    base64_image2 = base64.b64encode(response2.content).decode('utf-8')

    user_content_analysis = [
        {
            "type": "text",
            "text": """These are the winner creative image for the auto insurance ad.
Analyse these images and list down elements for these images"""
        },
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{base64_image1}"}
        },
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{base64_image2}"}
        },
    ]
    resp_analyses = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": user_content_analysis}]
    )
    analyses = resp_analyses.choices[0].message.content

    user_content_story = [
        {
            "type": "text",
            "text": """Based on this analysis, generate me the new image story for the auto insurance ad 
which should have new offers, revised text and graphics but style should remain same.
Make sure you only add text if it is present in previous images"""
        }
    ]
    resp_story = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": user_content_analysis
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": analyses}]
            },
            {
                "role": "user",
                "content": user_content_story
            },
        ]
    )
    story = resp_story.choices[0].message.content

    final_user_content_prompt = [
        {
            "type": "text",
            "text": """Now convert this image description to the prompt for the image generation through ideogram model.
The prompt should be short and crisp for it to understand really well.
Please also consider the background colors, text colors, text overlays and any other elements in the image to be included in prompt.
Make sure you only return the prompt.
"""
        }
    ]
    resp_prompt = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": user_content_analysis
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": analyses}]
            },
            {
                "role": "user",
                "content": user_content_story
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": story}]
            },
            {
                "role": "user",
                "content": final_user_content_prompt
            },
        ]
    )
    prompt = resp_prompt.choices[0].message.content
    return prompt


def generate_gpt_prompt_advanced_feature(results, sentiment, image_elements=None, text_elements=None, input_image=None, is_text=True, non_compliant=None, emotion=None):

    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    response1 = requests.get(results[0]['image_url'])
    base64_image1 = base64.b64encode(response1.content).decode('utf-8')

    response2 = requests.get(results[1]['image_url'])
    base64_image2 = base64.b64encode(response2.content).decode('utf-8')

    story_prompt = f"Based on this analysis, generate me the new image story for the auto insurance ad whose Sentiment of the new image story MUST be {sentiment}."

    if image_elements:
        story_prompt += f"You MUST creatively add these {image_elements} image elements."

    if is_text is True:
        story_prompt += f"Also, MUST add these '{text_elements}' text in new image story."
    else:
        story_prompt += "DON'T add any text overlays on the image."

    if non_compliant:
        story_prompt += f"DON'T add the mentioned stuff in the image story {non_compliant}."

    if emotion:
        story_prompt += f"The image should reflect {emotion} emotion."

    story_prompt += "Blend everything with new image story which should have new offers, revised text and graphics but style should remain same."


    user_content_analysis = [
        {
            "type": "text",
            "text": """These are the winner creative image for the auto insurance ad.
Analyse these images and list down elements for these images"""
        },
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{base64_image1}"}
        },
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{base64_image2}"}
        },
    ]
    resp_analyses = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": user_content_analysis}]
    )
    analyses = resp_analyses.choices[0].message.content

    if input_image:
        input_image_base = _encode_image_to_base64(input_image)
        story_prompt += "Also, consider the reference image while creating the new story."
        user_content_story = [
            {
                "type": "text",
                "text": story_prompt
            },
            {
                "type": "image_url",
                "image_url": f"data:image/jpeg;base64,{input_image_base}"
            }
        ]
    else:
        user_content_story = [
            {
                "type": "text",
                "text": story_prompt
            }
        ]
    # print(story_prompt)
    resp_story = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": user_content_analysis
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": analyses}]
            },
            {
                "role": "user",
                "content": user_content_story
            },
        ]
    )
    story = resp_story.choices[0].message.content

    final_user_content_prompt = [
        {
            "type": "text",
            "text": """Now convert this image description to the prompt for the image generation through ideogram model.
The prompt should be short and crisp for it to understand really well.
Please also consider the background colors, text colors, text overlays and any other elements in the image to be included in prompt.
Make sure you only return the prompt.
"""
        }
    ]
    resp_prompt = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": user_content_analysis
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": analyses}]
            },
            {
                "role": "user",
                "content": user_content_story
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": story}]
            },
            {
                "role": "user",
                "content": final_user_content_prompt
            },
        ]
    )
    prompt = resp_prompt.choices[0].message.content

    print(prompt)
    return prompt