from langchain.prompts import PromptTemplate
from typing import Dict, List
import json
import requests
import os

# For Groq
class GroqLLM:
    def __init__(self, api_key: str, model: str = "openai/gpt-oss-120b"):
        self.api_key = api_key
        self.model = model
        self.base_url = "https://api.groq.com/openai/v1"
    
    def __call__(self, prompt: str) -> str:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.1,
            "max_tokens": 1024
        }
        
        response = requests.post(
            f"{self.base_url}/chat/completions",
            headers=headers,
            json=data
        )
        
        if response.status_code == 200:
            return response.json()["choices"][0]["message"]["content"]
        else:
            raise Exception(f"Groq API error: {response.text}")


class FreeChatService:
    def __init__(self, llm_provider: str, **kwargs):
        self.llm_provider = llm_provider
        
        if llm_provider == "groq":
            self.llm = GroqLLM(
                api_key=kwargs.get("api_key"),
                model=kwargs.get("model", "openai/gpt-oss-120b")
            )
        else:
            raise ValueError(f"Unsupported LLM provider: {llm_provider}. Only 'groq' is supported.")
        
        self.prompt_template = PromptTemplate(
            input_variables=["context", "question"],
            template="""
You are a helpful AI assistant that analyzes code repositories. Use the following code snippets to answer the user's question about the repository.

Context from repository:
{context}

Question: {question}

Please provide a detailed answer based on the code context provided. If you reference specific files or functions, mention their file paths. If the question cannot be fully answered from the provided context, say so clearly.

Answer:"""
        )
    
    async def answer_question(self, question: str, vectorstore, repo_id: str) -> Dict:
        """Answer question using RAG with Groq LLM"""
        try:
            # Manual RAG for Groq
            docs = vectorstore.similarity_search(question, k=5)
            context = "\n\n".join([doc.page_content for doc in docs])
            
            prompt = self.prompt_template.format(
                context=context,
                question=question
            )
            
            answer = self.llm(prompt)
            source_docs = docs
            
            # Format sources
            sources = []
            for doc in source_docs:
                sources.append({
                    "path": doc.metadata.get("path", "Unknown"),
                    "content_preview": doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content
                })
            
            return {
                "response": answer,
                "sources": sources,
                "repo_id": repo_id
            }
        
        except Exception as e:
            return {
                "response": f"Error processing question: {str(e)}",
                "sources": [],
                "repo_id": repo_id
            }