modcomms/backend/app/services/gemini_service.py

import json
import logging
from google import genai
from google.genai import types

from app.models.schemas import SubReview, RagStatus

# Configure logging
logger = logging.getLogger(__name__)


class GeminiService:
    """Service wrapper for Google Gemini API calls."""

    def __init__(self, api_key: str):
        """
        Initialize the Gemini service.

        Args:
            api_key: Google Gemini API key
        """
        self.client = genai.Client(api_key=api_key)
        self.model = "gemini-2.5-flash"

    async def analyze_with_image(
        self,
        prompt: str,
        file_data: bytes,
        file_type: str,
    ) -> SubReview:
        """
        Analyze an image/file with Gemini and return a structured SubReview.

        Args:
            prompt: The analysis prompt including reference doc context
            file_data: Raw bytes of the file
            file_type: MIME type of the file (e.g., "image/png")

        Returns:
            SubReview with ragStatus, feedback, and issues
        """
        try:
            logger.info(f"[GEMINI API] Starting image analysis - file_type: {file_type}, file_size: {len(file_data)} bytes")

            # Create inline data part for the file
            file_part = types.Part.from_bytes(
                data=file_data,
                mime_type=file_type
            )

            # Define the response schema for structured output
            response_schema = {
                "type": "object",
                "properties": {
                    "analysisStatus": {
                        "type": "string",
                        "enum": ["success", "low_confidence"],
                        "description": "Set to 'low_confidence' if the proof is nonsensical, completely irrelevant to marketing, or otherwise impossible to analyze. Otherwise, set to 'success'."
                    },
                    "ragStatus": {
                        "type": "string",
                        "enum": ["Red", "Amber", "Green"],
                        "description": "A RAG status. Red: Issues that must be resolved. Amber: Issues that should be addressed. Green: No issues found."
                    },
                    "feedback": {
                        "type": "string",
                        "description": "Constructive, professional feedback explaining the RAG status and highlighting both positive aspects and areas for improvement."
                    },
                    "issues": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "A list of specific, actionable issues found. If no issues, return an empty array."
                    }
                },
                "required": ["analysisStatus", "ragStatus", "feedback", "issues"]
            }

            # Make the API call
            logger.info(f"[GEMINI API] Calling Gemini model: {self.model}")
            response = await self.client.aio.models.generate_content(
                model=self.model,
                contents=[file_part, prompt],
                config=types.GenerateContentConfig(
                    response_mime_type="application/json",
                    response_schema=response_schema
                )
            )
            logger.info(f"[GEMINI API] Response received from Gemini")

            # Parse the JSON response
            json_text = response.text.strip()
            parsed_result = json.loads(json_text)
            logger.info(f"[GEMINI API] Parsed result - ragStatus: {parsed_result.get('ragStatus')}, analysisStatus: {parsed_result.get('analysisStatus')}")

            # Handle low confidence analysis
            if parsed_result.get("analysisStatus") == "low_confidence":
                return SubReview(
                    ragStatus=RagStatus.ERROR,
                    feedback="The agent could not analyze this proof with high confidence. This may be because the content is irrelevant, nonsensical, or too far outside of expected marketing materials.",
                    issues=[]
                )

            # Return successful analysis
            return SubReview(
                ragStatus=RagStatus(parsed_result["ragStatus"]),
                feedback=parsed_result["feedback"],
                issues=parsed_result["issues"]
            )

        except json.JSONDecodeError as e:
            logger.error(f"[GEMINI API] JSON parse error: {str(e)}")
            return SubReview(
                ragStatus=RagStatus.ERROR,
                feedback=f"Failed to parse AI response as JSON: {str(e)}",
                issues=[]
            )
        except Exception as e:
            logger.error(f"[GEMINI API] Error during analysis: {str(e)}")
            return SubReview(
                ragStatus=RagStatus.ERROR,
                feedback=f"An error occurred during analysis: {str(e)}",
                issues=[]
            )

    async def generate_summary(
        self,
        prompt: str,
    ) -> dict:
        """
        Generate a text summary (for lead agent).

        Args:
            prompt: The prompt for generating the summary

        Returns:
            Parsed JSON response with summary and overall_status
        """
        try:
            logger.info("[GEMINI API] Generating lead agent summary")
            response_schema = {
                "type": "object",
                "properties": {
                    "overallStatus": {
                        "type": "string",
                        "enum": ["Passed", "Failed", "Analysis Error", "Requires Manual Legal Review"],
                    },
                    "summary": {
                        "type": "string",
                        "description": "A concise, professional summary explaining the overall status, based on the specialist reviews."
                    }
                },
                "required": ["overallStatus", "summary"]
            }

            response = await self.client.aio.models.generate_content(
                model=self.model,
                contents=prompt,
                config=types.GenerateContentConfig(
                    response_mime_type="application/json",
                    response_schema=response_schema
                )
            )

            result = json.loads(response.text.strip())
            logger.info(f"[GEMINI API] Summary generated - overallStatus: {result.get('overallStatus')}")
            return result

        except Exception as e:
            logger.error(f"[GEMINI API] Error generating summary: {str(e)}")
            return {
                "overallStatus": "Analysis Error",
                "summary": f"An error occurred while generating the summary: {str(e)}"
            }