"""
Question Answer Helper Utility

World-class utility for answering all 5 question types in domain assessments.
Provides intelligent, reliable, and fast question answering logic.
"""
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException
import random
import re


class QuestionAnswerHelper:
    """Helper class for answering questions in domain assessments"""
    
    def __init__(self, driver):
        """
        Initialize Question Answer Helper
        
        Args:
            driver: WebDriver instance
        """
        self.driver = driver
        self.wait_timeout = 10
    
    def get_question_id(self, question_element=None):
        """
        Extract question ID from question element or current page
        
        Args:
            question_element: Optional WebElement for question container
            
        Returns:
            str: Question ID or None
        """
        try:
            if question_element:
                test_id = question_element.get_attribute("data-testid")
            else:
                # Find current question on page - wait for it to appear
                try:
                    question_element = WebDriverWait(self.driver, 10).until(
                        EC.presence_of_element_located((By.CSS_SELECTOR, "[data-testid^='domain_question__']"))
                    )
                    test_id = question_element.get_attribute("data-testid")
                except:
                    # Try finding by question container (without __)
                    try:
                        question_element = WebDriverWait(self.driver, 5).until(
                            EC.presence_of_element_located((By.CSS_SELECTOR, "[data-testid^='domain_question__'][data-testid$='']"))
                        )
                        test_id = question_element.get_attribute("data-testid")
                    except:
                        # Last resort: find any question element
                        question_elements = self.driver.find_elements(By.CSS_SELECTOR, "[data-testid^='domain_question__']")
                        if question_elements:
                            # Get the first one that matches the pattern (not a sub-element)
                            for elem in question_elements:
                                test_id = elem.get_attribute("data-testid")
                                if test_id and not any(x in test_id for x in ['__option_', '__truefalse_', '__rating_', '__textarea', '__matrix_', '__header', '__text', '__number']):
                                    question_element = elem
                                    break
            
            if question_element:
                test_id = question_element.get_attribute("data-testid")
                if test_id:
                    match = re.search(r'domain_question__(\d+)(?:__|$)', test_id)
                    if match:
                        return match.group(1)
        except Exception as e:
            print(f"⚠️  Error getting question ID: {e}")
        return None
    
    def get_question_type(self, question_id):
        """
        Determine question type by checking available answer elements
        
        Args:
            question_id: Question ID
            
        Returns:
            str: Question type (multiple_choice, true_false, rating_scale, open_ended, matrix)
        """
        try:
            # Wait a moment for question to fully render
            import time
            time.sleep(0.5)
            
            # First, check for container elements (more reliable)
            # Check for multiple choice container
            if self._element_exists(f"[data-testid='domain_question__{question_id}__multiple_choice']", timeout=3):
                return "multiple_choice"
            
            # Check for true/false container
            if self._element_exists(f"[data-testid='domain_question__{question_id}__true_false']", timeout=3):
                return "true_false"
            
            # Check for rating scale container
            if self._element_exists(f"[data-testid='domain_question__{question_id}__rating_scale']", timeout=3):
                return "rating_scale"
            
            # Check for open ended container
            if self._element_exists(f"[data-testid='domain_question__{question_id}__open_ended']", timeout=3):
                return "open_ended"
            
            # Check for matrix container
            if self._element_exists(f"[data-testid='domain_question__{question_id}__matrix']", timeout=3):
                return "matrix"
            
            # Fallback: Check for individual answer elements (if containers not found)
            # Check for multiple choice options (A, B, C, D, E)
            for label in ['A', 'B', 'C', 'D', 'E']:
                if self._element_exists(f"[data-testid='domain_question__{question_id}__option_{label}']", timeout=2):
                    return "multiple_choice"
            
            # Check for true/false buttons
            if self._element_exists(f"[data-testid='domain_question__{question_id}__truefalse_True']", timeout=2):
                return "true_false"
            if self._element_exists(f"[data-testid='domain_question__{question_id}__truefalse_False']", timeout=2):
                return "true_false"
            
            # Check for rating scale buttons
            for rating in ['1', '2', '3', '4', '5']:
                if self._element_exists(f"[data-testid='domain_question__{question_id}__rating_{rating}']", timeout=2):
                    return "rating_scale"
            
            # Check for open ended textarea
            if self._element_exists(f"[data-testid='domain_question__{question_id}__textarea']", timeout=2):
                return "open_ended"
            
            # Check for matrix cells
            if self._element_exists(f"[data-testid^='domain_question__{question_id}__matrix_']", timeout=2):
                return "matrix"
            
            return "unknown"
        except Exception as e:
            print(f"⚠️  Error detecting question type for {question_id}: {e}")
            return "unknown"
    
    def _element_exists(self, css_selector, timeout=2):
        """Check if element exists quickly"""
        try:
            WebDriverWait(self.driver, timeout).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, css_selector))
            )
            return True
        except:
            return False
    
    def answer_multiple_choice(self, question_id, option_label=None):
        """
        Answer multiple choice question
        
        Args:
            question_id: Question ID
            option_label: Option label (A, B, C, D, E) - if None, selects random
            
        Returns:
            str: Selected option label
        """
        # Get available options
        options = []
        for label in ['A', 'B', 'C', 'D', 'E']:
            locator = (By.CSS_SELECTOR, f"[data-testid='domain_question__{question_id}__option_{label}']")
            try:
                element = WebDriverWait(self.driver, 2).until(
                    EC.presence_of_element_located(locator)
                )
                if element.is_displayed():
                    options.append(label)
            except:
                continue
        
        if not options:
            raise Exception(f"No options found for question {question_id}")
        
        # Select option
        if option_label is None:
            option_label = random.choice(options)
        elif option_label not in options:
            option_label = random.choice(options)  # Fallback to random
        
        locator = (By.CSS_SELECTOR, f"[data-testid='domain_question__{question_id}__option_{option_label}']")
        element = WebDriverWait(self.driver, self.wait_timeout).until(
            EC.element_to_be_clickable(locator)
        )
        element.click()
        
        return option_label
    
    def answer_true_false(self, question_id, value=None):
        """
        Answer true/false question
        
        Args:
            question_id: Question ID
            value: True or False - if None, selects random
            
        Returns:
            str: Selected value ("True" or "False")
        """
        if value is None:
            value = random.choice([True, False])
        
        value_str = "True" if value else "False"
        locator = (By.CSS_SELECTOR, f"[data-testid='domain_question__{question_id}__truefalse_{value_str}']")
        
        element = WebDriverWait(self.driver, self.wait_timeout).until(
            EC.element_to_be_clickable(locator)
        )
        element.click()
        
        return value_str
    
    def answer_rating_scale(self, question_id, score=None):
        """
        Answer rating scale question
        
        Args:
            question_id: Question ID
            score: Rating score/value - if None, selects random
            
        Returns:
            str: Selected score/value
        """
        # First, find all rating options dynamically (not just '1'-'5')
        # Rating options can have any value (numeric strings, labels, etc.)
        rating_options = []
        
        # Method 1: Find all elements with rating pattern
        try:
            all_rating_elements = self.driver.find_elements(
                By.CSS_SELECTOR, 
                f"[data-testid^='domain_question__{question_id}__rating_']"
            )
            for elem in all_rating_elements:
                test_id = elem.get_attribute('data-testid')
                if test_id:
                    # Extract value from pattern: domain_question__{id}__rating_{value}
                    match = re.search(r'__rating_(.+)$', test_id)
                    if match:
                        value = match.group(1)
                        if elem.is_displayed():
                            rating_options.append(value)
        except:
            pass
        
        # Method 2: Fallback - try common numeric values if nothing found
        if not rating_options:
            for s in ['1', '2', '3', '4', '5']:
                locator = (By.CSS_SELECTOR, f"[data-testid='domain_question__{question_id}__rating_{s}']")
                try:
                    element = WebDriverWait(self.driver, 2).until(
                        EC.presence_of_element_located(locator)
                    )
                    if element.is_displayed():
                        rating_options.append(s)
                except:
                    continue
        
        if not rating_options:
            raise Exception(f"No rating options found for question {question_id}")
        
        # Select score
        if score is None:
            selected_value = random.choice(rating_options)
        else:
            score_str = str(score)
            if score_str in rating_options:
                selected_value = score_str
            else:
                # Fallback to random if provided score not found
                selected_value = random.choice(rating_options)
        
        locator = (By.CSS_SELECTOR, f"[data-testid='domain_question__{question_id}__rating_{selected_value}']")
        
        element = WebDriverWait(self.driver, self.wait_timeout).until(
            EC.element_to_be_clickable(locator)
        )
        element.click()
        
        return selected_value
    
    def answer_open_ended(self, question_id, text=None):
        """
        Answer open-ended question
        
        Args:
            question_id: Question ID
            text: Answer text - if None, generates default text
            
        Returns:
            str: Entered text
        """
        if text is None:
            text = "This is a thoughtful response to the open-ended question. I believe this answer demonstrates understanding and reflection."
        
        locator = (By.CSS_SELECTOR, f"[data-testid='domain_question__{question_id}__textarea']")
        
        element = WebDriverWait(self.driver, self.wait_timeout).until(
            EC.presence_of_element_located(locator)
        )
        element.clear()
        element.send_keys(text)
        
        return text
    
    def answer_matrix(self, question_id, row_index=None, column_index=None):
        """
        Answer matrix question
        
        Args:
            question_id: Question ID
            row_index: Row index (0-based) - if None, selects random
            column_index: Column index (0-based) - if None, selects random
            
        Returns:
            tuple: (row_index, column_index) selected
        """
        # Get matrix dimensions
        matrix_cells = self.driver.find_elements(
            By.CSS_SELECTOR, 
            f"[data-testid^='domain_question__{question_id}__matrix_']"
        )
        
        if not matrix_cells:
            raise Exception(f"No matrix cells found for question {question_id}")
        
        # Extract available row/column indices
        rows = set()
        cols = set()
        for cell in matrix_cells:
            test_id = cell.get_attribute("data-testid")
            if test_id:
                match = re.search(r'matrix_(\d+)_(\d+)', test_id)
                if match:
                    rows.add(int(match.group(1)))
                    cols.add(int(match.group(2)))
        
        if not rows or not cols:
            raise Exception(f"Could not determine matrix dimensions for question {question_id}")
        
        # Select indices
        if row_index is None:
            row_index = random.choice(list(rows))
        elif row_index not in rows:
            row_index = random.choice(list(rows))  # Fallback
        
        if column_index is None:
            column_index = random.choice(list(cols))
        elif column_index not in cols:
            column_index = random.choice(list(cols))  # Fallback
        
        # Click matrix cell
        locator = (By.CSS_SELECTOR, 
                  f"[data-testid='domain_question__{question_id}__matrix_{row_index}_{column_index}']")
        
        element = WebDriverWait(self.driver, self.wait_timeout).until(
            EC.element_to_be_clickable(locator)
        )
        element.click()
        
        return (row_index, column_index)
    
    def answer_question(self, question_id=None, question_type=None, **kwargs):
        """
        Universal method to answer any question type
        
        Args:
            question_id: Question ID (if None, auto-detects)
            question_type: Question type (if None, auto-detects)
            **kwargs: Additional arguments for specific question types
                - option_label: For multiple_choice
                - value: For true_false
                - score: For rating_scale
                - text: For open_ended
                - row_index, column_index: For matrix
        
        Returns:
            dict: Answer details
        """
        # Auto-detect question ID if not provided
        if question_id is None:
            question_id = self.get_question_id()
            if not question_id:
                raise Exception("Could not determine question ID")
        
        # Auto-detect question type if not provided
        if question_type is None:
            question_type = self.get_question_type(question_id)
            if question_type == "unknown":
                raise Exception(f"Could not determine question type for question {question_id}")
        
        # Answer based on type
        result = {
            'question_id': question_id,
            'question_type': question_type,
            'answer': None
        }
        
        if question_type == "multiple_choice":
            option_label = kwargs.get('option_label')
            result['answer'] = self.answer_multiple_choice(question_id, option_label)
            result['option_label'] = result['answer']
        
        elif question_type == "true_false":
            value = kwargs.get('value')
            result['answer'] = self.answer_true_false(question_id, value)
            result['value'] = result['answer']
        
        elif question_type == "rating_scale":
            score = kwargs.get('score')
            result['answer'] = self.answer_rating_scale(question_id, score)
            result['score'] = result['answer']
        
        elif question_type == "open_ended":
            text = kwargs.get('text')
            result['answer'] = self.answer_open_ended(question_id, text)
            result['text'] = result['answer']
        
        elif question_type == "matrix":
            row_index = kwargs.get('row_index')
            column_index = kwargs.get('column_index')
            row_idx, col_idx = self.answer_matrix(question_id, row_index, column_index)
            result['answer'] = f"matrix_{row_idx}_{col_idx}"
            result['row_index'] = row_idx
            result['column_index'] = col_idx
        
        else:
            raise Exception(f"Unsupported question type: {question_type}")
        
        return result

