from sentence_transformers import SentenceTransformer
from faiss import IndexFlatIP
import numpy as np
from huggingface_hub import InferenceClient
import json
import os

class RAGInterviewAgent:
    def __init__(self, job_role, cv_summary):
        self.job_role = job_role
        self.cv_summary = cv_summary
        self.current_q = 0
        
        # Initialize models
        self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
        self.llm = InferenceClient(model="meta-llama/Meta-Llama-3-8B-Instruct")
        
        # Setup FAISS index
        self.index = IndexFlatIP(384)  # Embedding dimension
        self._setup_knowledge_base()
        
        # Generate questions
        self.questions = self._generate_questions()
    
    def _setup_knowledge_base(self):
        """Load job requirements into vector store"""
        reqs = {
            "Software Engineer": [
                "Explain SOLID principles",
                "How would you optimize a slow SQL query?",
                "Describe your experience with unit testing"
            ],
            "Data Scientist": [
                "Explain bias-variance tradeoff",
                "How would you handle missing data?",
                "Describe a machine learning project you worked on"
            ]
        }
        
        self.knowledge = reqs[self.job_role]
        embeddings = self.embedder.encode(self.knowledge)
        self.index.add(embeddings)
    
    def _generate_questions(self):
        """Generate questions using RAG"""
        questions = []
        
        # Get most relevant requirements
        cv_embed = self.embedder.encode(self.cv_summary["text"])
        _, indices = self.index.search(np.array([cv_embed]), 3)
        
        for i in indices[0]:
            prompt = f"""
            Generate an interview question for a {self.job_role} position
            based on this requirement: {self.knowledge[i]}
            Candidate's CV summary: {self.cv_summary['text'][:1000]}
            """
            
            response = self.llm.text_generation(
                prompt,
                max_new_tokens=100,
                temperature=0.7
            )
            questions.append({
                "text": response.strip(),
                "base_requirement": self.knowledge[i]
            })
        
        return questions
    
    def get_current_question(self):
        return self.questions[self.current_q]
    
    def next_question(self):
        self.current_q += 1
    
    def evaluate_answer(self, answer):
        """Evaluate answer against knowledge base"""
        # Get relevant requirement
        req = self.questions[self.current_q]["base_requirement"]
        
        prompt = f"""
        Evaluate this interview answer on a scale of 1-10:
        Requirement: {req}
        Answer: {answer}
        
        Provide JSON output with: score, feedback
        """
        
        response = self.llm.text_generation(
            prompt,
            max_new_tokens=200,
            temperature=0.3
        )
        
        try:
            return json.loads(response.split("{")[1].split("}")[0] + "}")
        except:
            return {"score": 5, "feedback": "Evaluation failed"}