# RAG Bot Core Implementation
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential
from openai import AzureOpenAI
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import logging

class QueryRequest(BaseModel):
    question: str
    context: str = ""

class RAGBot:
    def __init__(self, search_endpoint, search_key, openai_endpoint, openai_key):
        self.search_client = SearchClient(
            endpoint=search_endpoint,
            index_name="documentation-index",
            credential=AzureKeyCredential(search_key)
        )
        self.openai_client = AzureOpenAI(
            azure_endpoint=openai_endpoint,
            api_key=openai_key,
            api_version="2024-02-15-preview"
        )
    
    def search_documents(self, query: str, top_k: int = 5):
        """Search for relevant documentation chunks"""
        try:
            results = self.search_client.search(
                query,
                top=top_k,
                select=["content", "source", "title", "url"],
                include_total_count=True
            )
            return list(results)
        except Exception as e:
            logging.error(f"Search error: {e}")
            return []
    
    def generate_answer(self, question: str, context_docs: list):
        """Generate answer using retrieved context"""
        if not context_docs:
            return "I couldn't find relevant documentation to answer your question."
        
        # Prepare context from retrieved documents
        context_text = "\n\n".join([
            f"Source: {doc['source']}\n{doc['content']}"
            for doc in context_docs
        ])
        
        # Create prompt for GPT-5
        system_prompt = """You are a helpful documentation assistant. Use the provided 
        documentation context to answer questions accurately. Always cite your sources 
        and provide relevant links when available."""
        
        user_prompt = f"""
        Question: {question}
        
        Documentation Context:
        {context_text}
        
        Please provide a helpful answer based on the documentation above:
        """
        
        try:
            response = self.openai_client.chat.completions.create(
                model="gpt-5",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,
                max_tokens=1000
            )
            
            return {
                "answer": response.choices[0].message.content,
                "sources": [doc['source'] for doc in context_docs],
                "confidence": self.calculate_confidence(context_docs)
            }
        except Exception as e:
            logging.error(f"OpenAI error: {e}")
            raise HTTPException(status_code=500, detail="Error generating answer")
    
    def calculate_confidence(self, context_docs: list):
        """Calculate confidence score based on search results"""
        if not context_docs:
            return 0.0
        
        # Simple confidence based on number and relevance of results
        avg_score = sum(doc.get('@search.score', 0) for doc in context_docs) / len(context_docs)
        return min(avg_score / 10.0, 1.0)  # Normalize to 0-1 range

# FastAPI application
app = FastAPI(title="RAG Documentation Bot")
rag_bot = RAGBot(
    search_endpoint=os.getenv("AZURE_SEARCH_ENDPOINT"),
    search_key=os.getenv("AZURE_SEARCH_KEY"),
    openai_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    openai_key=os.getenv("AZURE_OPENAI_KEY")
)

@app.post("/ask")
async def ask_question(request: QueryRequest):
    """Main endpoint for asking questions"""
    try:
        # Search for relevant documents
        context_docs = rag_bot.search_documents(request.question)
        
        # Generate answer
        result = rag_bot.generate_answer(request.question, context_docs)
        
        return {
            "question": request.question,
            "answer": result["answer"],
            "sources": result["sources"],
            "confidence": result["confidence"]
        }
    except Exception as e:
        logging.error(f"Error processing question: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")