from functools import lru_cache
from typing import Any, Dict, List, Optional

import numpy as np
import spacy
import torch
from sentence_transformers import SentenceTransformer
from spacy.language import Language
from transformers import (
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    BlipForConditionalGeneration,
    BlipProcessor,
)


class QuestionGenerator:  # pylint: disable=too-many-instance-attributes
    """Handles question generation from text and images."""

    def __init__(self, nlp: Optional[Language] = None, sentence_model=None):
        """Initialize the question generator with models."""
        # Initialize model configuration first
        self.qa_model_name = "distilbert-base-cased-distilled-squad"

        # Initialize tokenizer and model
        self.qa_tokenizer = self._load_qa_tokenizer()
        self.qa_model = self._load_qa_model()
        self.blip_processor = self._load_blip_processor()
        self.blip_model = self._load_blip_model()

        # Allow injection for testing
        if nlp is not None:
            self.nlp = nlp
        else:
            try:
                self.nlp = spacy.load(
                    "en_core_web_sm", exclude=["vectors", "textcat", "parser", "ner"]
                )
            except BaseException:  # Catch without binding to unused variable
                self.nlp = spacy.blank("en")
                self.nlp.add_pipe("sentencizer")

        self.sentence_model = sentence_model or SentenceTransformer("all-MiniLM-L6-v2")
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Move models to appropriate device
        self.qa_model.to(self.device)
        self.blip_model.to(self.device)

    @lru_cache(maxsize=1)  # Cache model loading
    def _load_qa_tokenizer(self):
        return AutoTokenizer.from_pretrained(self.qa_model_name)

    @lru_cache(maxsize=1)  # Cache model loading
    def _load_qa_model(self):
        return AutoModelForQuestionAnswering.from_pretrained(self.qa_model_name)

    @lru_cache(maxsize=1)  # Cache model loading
    def _load_blip_processor(self):
        return BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")

    @lru_cache(maxsize=1)  # Cache model loading
    def _load_blip_model(self):
        return BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

    def generate_question_from_text(self, context: str) -> Dict[str, Any]:
        """Generate a question from a given text context."""
        doc = self.nlp(context)
        sentences = list(doc.sents)
        if not sentences:
            raise ValueError("No sentences found in text")

        sentence = np.random.choice(sentences)
        question = f"Which of the following is correct about {sentence.root}?"
        answer = str(sentence)

        return {"text": question, "context": context, "answer": answer}

    def generate_distractors(
        self, correct_answer: str, context: str, num_distractors: int = 3
    ) -> List[str]:
        """Generate distractor options using semantic similarity."""
        if not correct_answer or not context:
            return ["Option A", "Option B", "Option C"]  # Default fallback

        # Extract sentences from context
        try:
            doc = self.nlp(context)
            sentences = [sent.text.strip() for sent in doc.sents if sent.text.strip()]

            if not sentences:
                return ["Option A", "Option B", "Option C"]

            # Convert sentences to numpy array for encoding
            sentences_array = np.array(sentences)
            sentence_embeddings = self.sentence_model.encode(sentences_array.tolist())
            correct_embedding = self.sentence_model.encode([correct_answer])[0]

            # Calculate similarities
            similarities = np.dot(sentence_embeddings, correct_embedding)

            # Get most dissimilar sentences as distractors
            sorted_indices = np.argsort(similarities)
            distractors = []

            for idx in sorted_indices:
                if len(distractors) >= num_distractors:
                    break
                if sentences[idx] != correct_answer:
                    distractors.append(sentences[idx])

            # If we don't have enough distractors, add some default ones
            while len(distractors) < num_distractors:
                distractors.append(f"Option {len(distractors) + 1}")

            return distractors[:num_distractors]

        except Exception:
            return ["Option A", "Option B", "Option C"]

    def create_quiz_question(self, context: str) -> Dict[str, Any]:
        """Create a complete quiz question with distractors."""
        question_data = self.generate_question_from_text(context)
        distractors = self.generate_distractors(question_data["answer"], context)
        options = [question_data["answer"]] + distractors
        np.random.shuffle(options)

        correct_index = options.index(question_data["answer"])

        return {
            "text": question_data["text"],
            "options": options,
            "correct_answer": correct_index,
            "context": context,
            "difficulty": "medium",
            "category": "auto-generated",
        }
