import pandas as pd
import ast
from groq import Groq
import os
from .prompts import captioning_prompt
from langchain_core.messages import FunctionMessage , AIMessage
from .tools import  InfluencerRetrievalTool
import re

class ImageCaptioner:
    def __init__(self, api_key=os.environ.get('GROQ_API_KEY')):
        self.client = Groq(api_key=api_key)

    def caption_image(self,image_base64,user_input):
        if len(image_base64)>0:
            print('Captioning image')
            chat_completion = self.client.chat.completions.create(
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": captioning_prompt(user_input)},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpg;base64,{image_base64[-1]}",
                                },
                            },
                        ],
                    }
                ],
                model="meta-llama/llama-4-scout-17b-16e-instruct",
                max_completion_tokens=50,
                temperature = 1
            )
            response=chat_completion.choices[0].message.content
            return response
        else:
            return ''
        
class AnalyticsViewer:
    def __init__(self, business_details):
        self.business_details = business_details
        
    def show_analytics(self):
        tool_response = InfluencerRetrievalTool().retrieve_for_analytics(str(self.business_details))
        return tool_response


class ResponseBlockExtractor:
    def __init__(self, response):
        self.response = response

    def extract_latest(self):
        latest_block = []
        temp_block = []

        # Reverse iterate through the messages
        for message in reversed(self.response):
            if isinstance(message, (FunctionMessage, AIMessage)):
                temp_block.insert(0, message.content)

                # Once we collect 3 items in correct structure, stop
                if len(temp_block) == 3:
                    if "tool=" in temp_block[1] and "query_response" in temp_block[1]:
                        latest_block = temp_block
                        break
                    else:
                        temp_block = []
        print('The latest block', latest_block)
        return latest_block


def handle_tools(tools_list,stored_data):
    tools_order = [
    "analytics",
    "ideation",
    "human-idea-refining",
    "generate-story",
    "generate-ultimate-story",
    "generate-image",
    ]
    if 'generate-story' in tools_list and len(stored_data['human_ideation_interactions'])<1:
        tools_list.append('human-idea-refining')
    if 'generate-ultimate-story' in tools_list and len(stored_data['brainstorming_response'])<1:
        tools_list.append('generate-story')
    
    tools_list = [tool for tool in tools_order if tool in tools_list]
    return tools_list