# src/llm_client.py

from abc import ABC, abstractmethod
from typing import Optional
import os
from dotenv import load_dotenv
from langsmith import traceable, get_current_run_tree

load_dotenv()


class LLMClient(ABC):
    """Abstract base class for LLM clients. LLM değişimi için interface."""

    @abstractmethod
    def generate(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.3,
        max_tokens: int = 2048
    ) -> str:
        pass


class GeminiClient(LLMClient):
    """Google Gemini LLM client (google-genai SDK)"""

    def __init__(
        self,
        model_name: str = "gemini-2.5-flash",
        api_key: Optional[str] = None
    ):
        try:
            from google import genai
        except ImportError:
            raise ImportError(
                "google-genai paketi kurulu değil. "
                "Kurmak için: pip install google-genai"
            )

        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "GEMINI_API_KEY bulunamadı. "
                ".env dosyasına GEMINI_API_KEY=xxx ekle."
            )

        self.model_name = model_name
        self._client = genai.Client(
            api_key=self.api_key,
            http_options={"api_version": "v1beta"}
        )

        print(f"Gemini client initialized: {model_name}")

    @traceable(name="Gemini Report Generation", run_type="llm",
               metadata={"model": "gemini-2.5-flash", "temperature": 0.3})
    def generate(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.3,
        max_tokens: int = 2048,
        min_response_chars: int = 800,
        max_retries: int = 2
    ) -> str:
        from google.genai import types

        full_prompt = f"{system_prompt}\n\n{user_prompt}"

        for attempt in range(max_retries + 1):
            try:
                response = self._client.models.generate_content(
                    model=self.model_name,
                    contents=full_prompt,
                    config=types.GenerateContentConfig(
                        temperature=temperature,
                        max_output_tokens=max_tokens,
                        thinking_config=types.ThinkingConfig(
                            thinking_budget=0
                        ),
                    )
                )
                text = response.text
                if len(text) >= min_response_chars or attempt == max_retries:
                    rt = get_current_run_tree()
                    if rt and response.usage_metadata:
                        um = response.usage_metadata
                        rt.add_metadata({
                            "usage": {
                                "input_tokens": um.prompt_token_count,
                                "output_tokens": um.candidates_token_count,
                                "total_tokens": um.total_token_count,
                            }
                        })
                    return text
                print(f"  → Short response ({len(text)} chars), retrying ({attempt + 1}/{max_retries})...")
            except Exception as e:
                if attempt == max_retries:
                    raise RuntimeError(f"Gemini generation failed: {e}")
                print(f"  → Generation error, retrying ({attempt + 1}/{max_retries}): {e}")

        raise RuntimeError("Gemini generation failed after retries")


# Test
if __name__ == "__main__":
    client = GeminiClient()

    response = client.generate(
        system_prompt="You are a helpful medical AI assistant.",
        user_prompt="Briefly explain what a kidney stone is in 2 sentences.",
        temperature=0.3
    )

    print("\nGemini Response:")
    print(response)