app/services/llm.py

import json
import logging
from google import genai
from google.genai import types
from app.config import settings

logger = logging.getLogger(__name__)


class GeminiService:
    def __init__(self):
        self._client = None

    def initialize(self):
        if not settings.gemini_api_key:
            logger.warning("GEMINI_API_KEY not set — LLM calls will fail")
            return
        self._client = genai.Client(api_key=settings.gemini_api_key)
        logger.info(f"Gemini client initialized with model: {settings.gemini_model}")

    async def generate(self, system_prompt: str, user_prompt: str) -> str:
        if not self._client:
            raise RuntimeError("Gemini client not initialized")

        response = self._client.models.generate_content(
            model=settings.gemini_model,
            contents=[types.Content(role="user", parts=[types.Part(text=user_prompt)])],
            config=types.GenerateContentConfig(
                system_instruction=system_prompt,
                thinking_config=types.ThinkingConfig(thinking_level="HIGH"),
            ),
        )
        # Extract text, skipping thinking parts
        text_parts = []
        for part in response.candidates[0].content.parts:
            if part.text and not getattr(part, "thought", False):
                text_parts.append(part.text)
        return "\n".join(text_parts).strip()

    async def generate_json(self, system_prompt: str, user_prompt: str) -> dict | list:
        raw = await self.generate(system_prompt, user_prompt)
        # Strip markdown code fences if present
        cleaned = raw.strip()
        if cleaned.startswith("```"):
            lines = cleaned.split("\n")
            # Remove first line (```json) and last line (```)
            lines = [l for l in lines[1:] if l.strip() != "```"]
            cleaned = "\n".join(lines)
        return json.loads(cleaned)


gemini_service = GeminiService()