app/core/model_loader.py

import os
import json
import logging
import threading
import time
import random
from datetime import datetime
from dotenv import load_dotenv
from groq import Groq
import google.generativeai as genai

load_dotenv()
logger = logging.getLogger(__name__)

STATS_FILE = "usage_stats.json"

# ====================== CONFIG ======================
# Gemini (Primary)
DEFAULT_GEMINI_MODEL = "gemini-2.5-pro"          # Higher quality. Use "gemini-2.5-flash" if you want faster + cheaper
FALLBACK_GEMINI_MODEL = "gemini-2.0-flash"

# Groq (Fallback)
DEFAULT_GROQ_MODEL = "openai/gpt-oss-120b"
FALLBACK_GROQ_MODEL = "openai/gpt-oss-120b"
# ====================================================


class LLMSingleton:
    _instance = None
    _instance_lock = threading.Lock()

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            with cls._instance_lock:
                if cls._instance is None:
                    cls._instance = cls()
        return cls._instance

    def __init__(self):
        if self._instance is not None:
            raise Exception("Singleton instance already exists!")

        # ---------- Gemini ----------
        self.gemini_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY") or ""
        self.gemini_key = self.gemini_key.strip().strip('"').strip("'")
        self.gemini_model_name = os.getenv("GEMINI_MODEL", DEFAULT_GEMINI_MODEL).strip() or DEFAULT_GEMINI_MODEL

        if self.gemini_key:
            try:
                genai.configure(api_key=self.gemini_key)
                self.gemini_model = genai.GenerativeModel(self.gemini_model_name)
                logger.info(f"✅ Gemini initialized → {self.gemini_model_name}")
            except Exception as e:
                logger.error(f"❌ Failed to init Gemini: {e}")
                self.gemini_model = None
        else:
            logger.warning("⚠️ GOOGLE_API_KEY / GEMINI_API_KEY not found")
            self.gemini_model = None

        # ---------- Groq (Fallback) ----------
        raw_key = os.getenv("GROQ_API_KEY", "")
        self.groq_key = raw_key.strip().strip('"').strip("'") if raw_key else ""
        self.groq_model_name = os.getenv("GROQ_MODEL", DEFAULT_GROQ_MODEL).strip() or DEFAULT_GROQ_MODEL

        if self.groq_key:
            self.groq_client = Groq(api_key=self.groq_key)
            logger.info(f"✅ Groq initialized → {self.groq_model_name}")
        else:
            logger.warning("⚠️ GROQ_API_KEY not found")
            self.groq_client = None

        # ---------- Stats & Rate limiting ----------
        self._stats_lock = threading.Lock()
        self._rpm_lock = threading.Lock()
        self.stats = self._load_stats()
        self._check_daily_reset()
        self.rpm_limit = 30
        self.minute_window_start = time.time()
        self.requests_this_minute = 0

    # ------------------------------------------------------------------
    # Stats helpers
    # ------------------------------------------------------------------
    def _load_stats(self):
        default_stats = {
            "total_requests": 0,
            "successful_requests": 0,
            "rate_limit_hits": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "errors": 0,
            "local_model_requests": 0,
            "date": datetime.now().strftime("%Y-%m-%d"),
            "daily_requests_count": 0,
        }
        if os.path.exists(STATS_FILE):
            try:
                with open(STATS_FILE, "r") as f:
                    data = json.load(f)
                    return {**default_stats, **data}
            except Exception as e:
                logger.error(f"Failed to load stats: {e}")
        return default_stats

    def _save_stats(self):
        try:
            with open(STATS_FILE, "w") as f:
                json.dump(self.stats, f, indent=2)
        except Exception as e:
            logger.error(f"Failed to save stats: {e}")

    def _check_daily_reset(self):
        with self._stats_lock:
            today = datetime.now().strftime("%Y-%m-%d")
            if self.stats.get("date") != today:
                logger.info("📅 New day detected. Resetting daily AI quotas.")
                self.stats["date"] = today
                self.stats["daily_requests_count"] = 0
                self._save_stats()

    def _check_rpm_window(self):
        with self._rpm_lock:
            now = time.time()
            if now - self.minute_window_start >= 60:
                self.minute_window_start = now
                self.requests_this_minute = 0

    def get_usage_stats(self):
        self._check_daily_reset()
        self._check_rpm_window()
        with self._stats_lock:
            stats = self.stats.copy()
        with self._rpm_lock:
            requests_this_minute = self.requests_this_minute

        daily_limit = 1000
        stats["limits"] = {
            "requests_per_minute": self.rpm_limit,
            "requests_per_day": daily_limit,
        }
        stats["remaining_daily_requests"] = max(0, daily_limit - stats["daily_requests_count"])
        stats["remaining_rpm"] = max(0, self.rpm_limit - requests_this_minute)
        return stats

    def _reserve_request_slot(self) -> bool:
        with self._stats_lock:
            if self.stats["daily_requests_count"] >= 1000:
                return False
            self.stats["total_requests"] += 1
            self.stats["daily_requests_count"] += 1
            self._save_stats()

        with self._rpm_lock:
            self.requests_this_minute += 1
        return True

    # ------------------------------------------------------------------
    # Core generation methods
    # ------------------------------------------------------------------
    def _call_gemini(self, prompt: str, system_prompt: str, max_tokens: int = 2048, json_mode: bool = False) -> str:
        if not self.gemini_model:
            raise RuntimeError("Gemini not available")

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

        generation_config = {
            "max_output_tokens": max_tokens,
            "temperature": 0.2 if json_mode else 0.3,
        }
        if json_mode:
            generation_config["response_mime_type"] = "application/json"

        response = self.gemini_model.generate_content(
            full_prompt,
            generation_config=generation_config,
        )
        return (response.text or "").strip()

    def _call_groq(self, prompt: str, system_prompt: str, max_tokens: int = 2048, json_mode: bool = False) -> str:
        if not self.groq_client:
            raise RuntimeError("Groq not available")

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ]

        kwargs = {
            "model": self.groq_model_name,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": 0.2 if json_mode else 0.3,
        }
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}

        completion = self.groq_client.chat.completions.create(**kwargs)
        return (completion.choices[0].message.content or "").strip()

    def generate(self, prompt: str, max_tokens: int = 2048) -> str:
        """Used for structured JSON responses (code review etc.)"""
        self._check_daily_reset()
        self._check_rpm_window()

        if not self._reserve_request_slot():
            raise RuntimeError("QUOTA_EXCEEDED")

        system_prompt = (
            "You are a senior Android code reviewer. "
            "You MUST return a valid JSON object matching the requested schema strictly."
        )

        # Try Gemini first
        if self.gemini_model:
            try:
                logger.info(f"🤖 Generating with Gemini ({self.gemini_model_name})")
                result = self._call_gemini(prompt, system_prompt, max_tokens, json_mode=True)
                with self._stats_lock:
                    self.stats["successful_requests"] += 1
                    self.stats["output_tokens"] += len(result) // 4
                    self._save_stats()
                return result
            except Exception as e:
                logger.warning(f"Gemini failed → falling back to Groq: {e}")

        # Fallback to Groq
        if self.groq_client:
            try:
                logger.info(f"🤖 Generating with Groq ({self.groq_model_name})")
                result = self._call_groq(prompt, system_prompt, max_tokens, json_mode=True)
                with self._stats_lock:
                    self.stats["successful_requests"] += 1
                    self.stats["output_tokens"] += len(result) // 4
                    self._save_stats()
                return result
            except Exception as e:
                logger.error(f"Groq also failed: {e}")
                raise RuntimeError(f"GENERATION_FAILED: {e}")

        raise RuntimeError("No LLM provider available")

    def generate_text(self, prompt: str) -> str:
        """Used for normal chat / explanations"""
        self._check_daily_reset()
        self._check_rpm_window()

        if not self._reserve_request_slot():
            return "Error: Daily Quota Exceeded."

        system_prompt = "You are GitGud AI, an expert software architect."

        # Try Gemini first
        if self.gemini_model:
            try:
                logger.info(f"🤖 Chat with Gemini ({self.gemini_model_name})")
                result = self._call_gemini(prompt, system_prompt, max_tokens=2048, json_mode=False)
                with self._stats_lock:
                    self.stats["successful_requests"] += 1
                    self.stats["output_tokens"] += len(result) // 4
                    self._save_stats()
                return result
            except Exception as e:
                logger.warning(f"Gemini chat failed → falling back to Groq: {e}")

        # Fallback to Groq
        if self.groq_client:
            try:
                logger.info(f"🤖 Chat with Groq ({self.groq_model_name})")
                result = self._call_groq(prompt, system_prompt, max_tokens=2048, json_mode=False)
                with self._stats_lock:
                    self.stats["successful_requests"] += 1
                    self.stats["output_tokens"] += len(result) // 4
                    self._save_stats()
                return result
            except Exception as e:
                logger.error(f"Groq chat also failed: {e}")
                return f"Error generating content: {str(e)}"

        return "Error: No LLM provider available (check API keys)."


# Export global instance
llm_engine = LLMSingleton.get_instance()