import os
import asyncio
from groq import Groq
from sqlalchemy.ext.asyncio import AsyncSession
from src.models.prediction import Prediction
from src.utils.files_utils import encode_image
from dotenv import load_dotenv

load_dotenv()

class ImageComparator:
    def __init__(self, db: AsyncSession):
        self.db = db
        self.client = Groq(api_key=os.getenv("GROQ_API_KEY"))

    async def generate_comparison(self, base64_1: str, base64_2: str) -> str:
        try:
            response = await asyncio.to_thread(
                self.client.chat.completions.create,
                model="meta-llama/llama-4-maverick-17b-128e-instruct",
                messages=[
                    {"role": "system", "content": "You are a professional construction site analyst with expertise in identifying and assessing construction materials, equipment, and safety conditions."},
                    {"role": "user", "content": [
                        {"type": "text", "text": """Compare these two construction site images and provide ONLY:

1. A concise bulleted list of changes between Image 1 and Image 2, focusing on:
   - New materials/equipment that have appeared
   - Materials/equipment that have been removed
   - Materials/equipment that have been relocated
   - Changes in quantity of materials
   - Progress indicators (e.g., "scaffolding extended", "bricks used")
   - Any new safety concerns

2. A single summary sentence about the overall progress made between the two images.

Focus on these construction elements:
- Wooden logs
- Concrete Mixer
- Rebar
- Scaffolding
- Doors
- Piles of Bricks
- Metallic Shuttering

Keep your response brief, factual, and structured for easy conversion to JSON format. Do not include any introductory text, analysis methodology, or conclusions beyond what is requested."""},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_1}"}},
                        {"type": "text", "text": "Image 1 (Earlier timestamp)"},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_2}"}},
                        {"type": "text", "text": "Image 2 (Later timestamp)"},
                    ]}
                ],
                temperature=0.3,
                max_completion_tokens=1000
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"Error during comparison: {str(e)}"

    async def run_comparison(self, prediction_id: int, before_path: str, after_path: str):
        try:
            base64_1 = await encode_image(before_path)
            base64_2 = await encode_image(after_path)
            comparison_text = await self.generate_comparison(base64_1, base64_2)
            prediction = await self.db.get(Prediction, prediction_id)
            if prediction:
                prediction.text_results = {"comparison_report": comparison_text}
                prediction.status = "processing"
                await self.db.commit()
                await self.db.refresh(prediction)
            else:
                raise Exception("Prediction not found")
            return {
                "comparison_report": comparison_text,
                "status": prediction.status
            }
        except Exception as e:
            raise Exception(f"Failed to run the text comparison: {str(e)}")