from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
import requests
import json

class BaseLLM(ABC):
    """Abstract base class for language models"""

    @abstractmethod
    def generate(self, prompt: str, **kwargs) -> str:
        """Generate text based on prompt"""
        pass

    @abstractmethod
    def generate_with_context(self, query: str, context: List[str], **kwargs) -> str:
        """Generate text using query and retrieved context"""
        pass

class GeminiLLM(BaseLLM):
    """Google Gemini language model"""

    def __init__(self, api_key: str, model: str = "gemini-2.5-flash"):
        self.api_key = api_key
        self.model = model
        self.base_url = "https://generativelanguage.googleapis.com/v1beta/models"

    def generate(self, prompt: str, **kwargs) -> str:
        """Generate text using Gemini API"""
        url = f"{self.base_url}/{self.model}:generateContent"

        headers = {
            "Content-Type": "application/json",
        }

        data = {
            "contents": [{
                "parts": [{"text": prompt}]
            }]
        }

        # Add generation config if provided
        generation_config = {}
        if "max_tokens" in kwargs:
            generation_config["maxOutputTokens"] = kwargs["max_tokens"]
        if "temperature" in kwargs:
            generation_config["temperature"] = kwargs["temperature"]

        if generation_config:
            data["generationConfig"] = generation_config

        response = requests.post(
            f"{url}?key={self.api_key}",
            headers=headers,
            json=data
        )

        if response.status_code == 200:
            result = response.json()
            if "candidates" in result and len(result["candidates"]) > 0:
                return result["candidates"][0]["content"]["parts"][0]["text"]
            else:
                raise Exception("No response generated")
        else:
            raise Exception(f"LLM API error: {response.status_code} - {response.text}")

    def generate_with_context(self, query: str, context: List[str], **kwargs) -> str:
        """Generate answer using query and retrieved context"""
        # Create a prompt that includes the context
        context_text = "\n\n".join([f"Context {i+1}: {ctx}" for i, ctx in enumerate(context)])

        prompt = f"""You are a helpful educational assistant. Use the provided context to answer the user's question. If the context doesn't contain relevant information, say so clearly.

Context:
{context_text}

Question: {query}

Answer:"""

        return self.generate(prompt, **kwargs)

class SimpleLLM(BaseLLM):
    """Simple fallback LLM that creates basic responses"""

    def __init__(self):
        pass

    def generate(self, prompt: str, **kwargs) -> str:
        """Generate a simple response (for testing when API is not available)"""
        return f"This is a simple response to: {prompt[:100]}..."

    def generate_with_context(self, query: str, context: List[str], **kwargs) -> str:
        """Generate answer with context (simple fallback)"""
        if context:
            return f"Based on the provided context, here's information related to '{query}': {context[0][:200]}..."
        else:
            return f"I don't have enough context to answer: {query}"

def get_llm_model(config) -> BaseLLM:
    """Factory function to get the appropriate LLM"""
    try:
        return GeminiLLM(config.gemini_api_key, config.llm_model)
    except Exception as e:
        print(f"Warning: Could not initialize Gemini LLM: {e}")
        print("Falling back to simple LLM for testing...")
        return SimpleLLM()