import os


import gradio as gr
# from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
import numpy as np
import torch
from shared_resources import shared_resources
from phi.agent import Agent
from phi.tools.duckduckgo import DuckDuckGo
from phi.agent import Agent, RunResponse
from phi.model.huggingface import HuggingFaceChat



class StoryRecommendationApp:

    def __init__(self):
        # self.tokenizer = shared_resources.tokenizer
        # self.model = shared_resources.model
        self.device = shared_resources.device
        self.sentence_transformer = shared_resources.sentence_transformer
        self.data = shared_resources.data

    def search(self, query: str, k: int = 3):
        """Search for recommended videos based on user input."""
        embedded_query = self.sentence_transformer.encode(query)
        scores, retrieved_examples = self.data.get_nearest_examples("embeddings", embedded_query, k=k)

        if isinstance(retrieved_examples, np.ndarray):
            retrieved_examples = np.nan_to_num(retrieved_examples, nan=0)
        elif isinstance(retrieved_examples, list):
            retrieved_examples = [
                [0 if np.isnan(val) else val for val in example] if isinstance(example, list) else example
                for example in retrieved_examples
            ]
        return scores, retrieved_examples

    def compute_mean_and_predict(self, retrieved_examples):
        """Compute mean for likesCount, commentCount, and shareCount."""
        features = ["LikesCount", "commentCount", "shareCount"]
        predictions = {}

        for feature in features:
            values = np.array(retrieved_examples[feature])
            values = np.array([np.nan if v is None else v for v in values])
            values = np.nan_to_num(values, nan=0.0)

            mean_value = int(np.mean(values))
            predictions[f"predicted_{feature}"] = mean_value

        return predictions

    def generate_prompt(self, query: str):
        """Generate a prompt for video generation based on the input story."""
        input_text = f'''
        I want to summarize a story in exactly 3 sentences. No more than 3 nor less than 3. 
        But the sentences have to be good enough to use as a prompt for video generation, because I have to give those 3 sentences to the video generation model.
        For example: This prompt is about A heartwarming family reunion celebrating love and cherished memories in exactly 3 sentences.
        -A warm, heartfelt reunion in a cozy living room, where family members embrace each other after a long time apart, soft lighting enhances the intimate atmosphere, and laughter fills the air.
        -A close-up shot of a grandmother’s hands carefully arranging a family photo album, as the camera pans over old pictures, evoking cherished memories and a deep sense of love.
        -A final moment around the dinner table, family members sharing a meal together, toasts are made, and the soft glow of candles reflects the joy and connection between generations.

        So, I will provide you that story now. The story is:\n{query}

        Please remember, don't give any background descriptions. Just generate 3 sentences likewise the example above. Don't even give the starting text like: "Here are the 3 sentences that summarize the story:" or any other like this. Just give the answers in 3 sentences directly
        '''
        agent = Agent(
            model=HuggingFaceChat(
                id="meta-llama/Meta-Llama-3-8B-Instruct",
                max_tokens=4096,
            ),
            # tools=[DuckDuckGo()],
            markdown=True
        )

        # Get the response in a variable
        run: RunResponse = agent.run(input_text)
        generated_text=run.content

        sentences = [sentence.strip() for sentence in generated_text.split('.') if sentence]
        return '. '.join(sentences[-3:]) + ('.' if len(sentences) > 0 else '')

    def generate_story_and_recommendation(self, generated_response: str):
        """Generate story recommendations and predictions based on the user input."""
        scores, result = self.search(generated_response, 4)

        recommended_videos_text = ""
        predictions = {}

        if scores is not None and result is not None:
            recommendations = []
            for idx in range(len(result['url'])):
                recommendations.append(
                    f"Video {idx+1}: {result['url'][idx]}\nPlaycount: {int(result['playCount'][idx])}\n"
                )
            recommended_text = "\n\n".join(recommendations)
            recommended_influencer = f"\nYou can use these influencers for this type of video {str(result['username'][:3])}"

            predictions = self.compute_mean_and_predict(result)
            generated_prompt = self.generate_prompt(generated_response)

        return recommended_text + recommended_influencer, predictions, generated_prompt

    def format_predictions(self, predictions):
        """Format predictions for display."""
        if predictions:
            return "\n".join([f"{key}: {value}" for key, value in predictions.items()])
        else:
            return "No predictions available."

    def launch_interface(self):
        """Launch the Gradio interface."""
        interface=gr.Interface(
            fn=self.generate_story_and_recommendation,
            inputs=gr.Textbox(label="Enter your generated story.", lines=15),
            outputs=[
                gr.Textbox(label="Our Recommendations for you."),
                gr.Textbox(label="Predicted Metrics (Likes, Comments, Shares)", type="text"),
                gr.Textbox(label="Recommended Prompt for video generation:"),
            ],
            title="Video Story Generation and Recommendation",
            description="Enter a request for a video storyline, and get a detailed story along with recommended videos and predicted engagement metrics based on the same input."
        )
        return interface