import base64
import httpx
from pymongo import MongoClient
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage
import os
import re
import json
from dotenv import load_dotenv
load_dotenv()
MONGO_URI = os.getenv("MONGO_URI")
DB_NAME = os.getenv("DB_NAME")
COLLECTION_NAME = os.getenv("COLLECTION_NAME")
FLASH_API = os.getenv("FLASH_API")
mongo_client = MongoClient(MONGO_URI)
db = mongo_client[DB_NAME]
collection = db[COLLECTION_NAME]
collection2=db['about_company']
model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0, max_tokens=None, google_api_key=FLASH_API)

about_company_doc=collection2.find_one({"type":"about_company"})
if about_company_doc:
    about_company=about_company_doc.get('company_description','')

system_prompt_text = f"""Given is an image related to a company. Your task is to analyze the image, identify any text or notable visual elements, and provide a comprehensive, direct description of the image's contents, focusing on what it represents without abstract language or additional commentary. The response must be concise and focused, using only descriptive nouns and adjectives. If the image cannot be clearly described, respond with 'None.'
Company information is given below to understand the context.
- About Company: {about_company}
Ensure you give a valid JSON structure.
Expected Output format : {{"description":"String"}}
"""




def process_image_using_llm(image_url):
    try:
        # Download and encode the image
        image_data = base64.b64encode(httpx.get(image_url).content).decode("utf-8")

        # Create the message with a system prompt and image
        message = HumanMessage(
            content=[
                {"type": "text", "text": system_prompt_text},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}},
            ],
        )

        # Invoke the model to extract information from the image
        response = model.invoke([message])
        print(f"llm response for {image_url} is : {response}")
        # Use regex to extract JSON part of the response
        match = re.search(r"\{.*\}", response.content.strip())
        if match:
            json_data = match.group(0)  # Extract JSON-like content as a string
            data = json.loads(json_data)  # Load as JSON

            # Get the description from the JSON data
            description = data.get("description", "None").strip()

            # Check if the description is "None"
            if description == "None":
                # Update MongoDB with can_find_description as False
                collection.update_one(
                    {"object_url": image_url},
                    {"$set": {"can_find_description": False}}
                )
                print(f"Marked {image_url} as can_find_description: False")
                return False
            else:
                # Update MongoDB with the description and set can_find_description to True
                collection.update_one(
                    {"object_url": image_url},
                    {"$set": {"description": description, "can_find_description": True}}
                )
                print("Saved description to MongoDB")
                return True
        else:
            print(f"No valid JSON found in the response for {image_url}. Marking as can_find_description: False")
            collection.update_one(
                {"object_url": image_url},
                {"$set": {"can_find_description": False}}
            )

            return False
    except Exception as e:
        print(f"Error processing {image_url}: {e}")
        return False