src/streamlit_app.py

import streamlit as st
import pandas as pd
import os
import hashlib
import json
import tempfile
from huggingface_hub import HfApi, login, hf_hub_download
import pycountry
from babel import Locale

# Hugging Face Dataset configuration
# In HF Spaces, variables and secrets are available as environment variables
# For local development, we can also check st.secrets as a fallback
HF_DATASET_REPO = os.getenv("HF_DATASET_REPO")
HF_TOKEN = os.getenv("HF_TOKEN")

# Fallback to st.secrets for local development (if not found in environment)
if not HF_DATASET_REPO:
    try:
        HF_DATASET_REPO = st.secrets.get("HF_DATASET_REPO", "TransLegal/grading-answers")
    except Exception:
        HF_DATASET_REPO = "TransLegal/grading-answers"

# Available jurisdictions will be discovered dynamically from the repository

if not HF_TOKEN:
    try:
        HF_TOKEN = st.secrets.get("HF_TOKEN", None)
    except Exception:
        HF_TOKEN = None

# Require HF_TOKEN - fail fast if not configured
if not HF_TOKEN:
    st.error("❌ **Configuration Error**: HF_TOKEN is not set. Please configure it in Hugging Face Spaces settings (Variables and secrets).")
    st.stop()

@st.cache_resource
def get_hf_api():
    """Get cached Hugging Face API client - only initializes once per session"""
    try:
        login(token=HF_TOKEN)
        return HfApi(token=HF_TOKEN)
    except Exception as e:
        st.error(f"❌ **Error initializing Hugging Face API**: {str(e)}")
        st.stop()

# Initialize HF API - cached to avoid re-initialization on every rerun
hf_api = get_hf_api()

@st.cache_data
def discover_available_jurisdictions():
    """
    Discover available jurisdictions from the Hugging Face dataset repository.
    Returns a list of jurisdiction subdirectories that contain grading_template.parquet.
    """
    try:
        # List all files in the repository
        repo_files = hf_api.list_repo_files(
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )
        
        # Extract unique jurisdiction names from file paths
        # Files are in format: "{jurisdiction}/grading_template.parquet" or "{jurisdiction}/users/..."
        jurisdictions = set()
        for file_path in repo_files:
            # Check if this is a grading_template.parquet file in a jurisdiction subdirectory
            if file_path.endswith("grading_template.parquet"):
                # Extract jurisdiction from path (e.g., "en-us/grading_template.parquet" -> "en-us")
                parts = file_path.split("/")
                if len(parts) == 2 and parts[0] and parts[1] == "grading_template.parquet":
                    jurisdictions.add(parts[0])
        
        # Return sorted list of jurisdictions
        available_jurisdictions = sorted(list(jurisdictions))
        
        if not available_jurisdictions:
            st.warning("⚠️ **Warning**: No jurisdictions found in the repository. Please ensure the repository structure is correct.")
            return []
        
        return available_jurisdictions
    except Exception as e:
        st.error(f"❌ **Error discovering jurisdictions from Hugging Face Dataset**: {str(e)}")
        st.error(f"Please ensure the repository {HF_DATASET_REPO} is accessible and contains jurisdiction subdirectories.")
        # Return empty list as fallback
        return []

def get_jurisdiction_display_name(jurisdiction_code):
    """
    Convert jurisdiction code (e.g., 'hr-hr') to display name (e.g., 'Croatian-Croatia').
    
    Uses ISO 639-1 (language) and ISO 3166-1 (country) codes to generate human-readable names.
    
    Args:
        jurisdiction_code: String in format 'language-country' (e.g., 'hr-hr', 'en-us')
    
    Returns:
        Display name string (e.g., 'Croatian-Croatia') or original code if conversion fails
    """
    try:
        # Parse jurisdiction code (e.g., "hr-hr" -> language="hr", country="HR")
        parts = jurisdiction_code.lower().split('-')
        if len(parts) != 2:
            return jurisdiction_code  # Fallback to original if format is wrong
        
        language_code, country_code = parts[0], parts[1].upper()
        
        # Get language name using babel
        language_name = None
        try:
            # Try to parse locale (e.g., "hr_HR" or "en_US")
            locale_str = f"{language_code}_{country_code}"
            locale = Locale.parse(locale_str)
            language_name = locale.get_language_name('en')
            if language_name:
                language_name = language_name.title()
        except Exception:
            pass
        
        # Fallback: try pycountry for language
        if not language_name:
            try:
                lang = pycountry.languages.get(alpha_2=language_code)
                language_name = lang.name
            except Exception:
                language_name = language_code.upper()
        
        # Get country name using pycountry
        country_name = None
        try:
            country = pycountry.countries.get(alpha_2=country_code)
            country_name = country.name
        except Exception:
            country_name = country_code
        
        return f"{language_name}-{country_name}"
    except Exception:
        # Fallback to original code if anything goes wrong
        return jurisdiction_code

@st.cache_data
def get_jurisdiction_display_mapping(jurisdiction_codes):
    """
    Create a mapping from jurisdiction codes to display names.
    
    Args:
        jurisdiction_codes: List of jurisdiction code strings
    
    Returns:
        Dictionary mapping codes to display names
    """
    return {code: get_jurisdiction_display_name(code) for code in jurisdiction_codes}

@st.cache_data
def load_grading_template(jurisdiction):
    """Load grading template from Hugging Face Dataset for the specified jurisdiction"""
    # Validate jurisdiction is set
    available_jurisdictions = discover_available_jurisdictions()
    if not jurisdiction or jurisdiction not in available_jurisdictions:
        st.error(f"❌ **Error**: Jurisdiction is not set or invalid. Please select a valid jurisdiction.")
        if available_jurisdictions:
            st.error(f"Available jurisdictions: {', '.join(available_jurisdictions)}")
        else:
            st.error("No jurisdictions found in the repository.")
        st.stop()
    
    try:
        file_path = hf_hub_download(
            repo_id=HF_DATASET_REPO,
            filename=f"{jurisdiction}/grading_template.parquet",
            repo_type="dataset",
            token=HF_TOKEN
        )
        return pd.read_parquet(file_path)
    except Exception as e:
        st.error(f"❌ **Error loading grading template from Hugging Face Dataset**: {str(e)}")
        st.error(f"Please ensure the file `{jurisdiction}/grading_template.parquet` exists in the dataset repository: {HF_DATASET_REPO}")
        st.stop()

# Assessment options with descriptive labels
ASSESSMENT_OPTIONS = [
    "Perfect",
    "Mostly correct",
    "Noticeably flawed",
    "Seriously wrong",
    "Irrelevant / NA"
]

# Map assessment options to scores
ASSESSMENT_TO_SCORE = {
    "Perfect": "3",
    "Mostly correct": "2",
    "Noticeably flawed": "1",
    "Seriously wrong": "0",
    "Irrelevant / NA": "NA"
}

# Score explanations from annotation guide
SCORE_EXPLANATIONS = {
    "Perfect": "Only when truly flawless. The answer is legally accurate, well-stated, and appropriate for legal education. It correctly explains the legal principle, rule, or concept without any discernible errors or misleading statements.",
    "Mostly correct": "One very small issue or slightly awkward phrasing. The answer is generally accurate but contains minor inaccuracies, imprecise language, or could be more precise. The core legal content is correct, but there are small issues that could be improved for educational purposes.",
    "Noticeably flawed": "Clear error but main idea still ok. The answer contains significant errors that substantially affect its accuracy. While not completely wrong, there are important mistakes in legal reasoning, application of law, or factual statements that would confuse or mislead students.",
    "Seriously wrong": "Hallucination, wrong format, meaningless, etc. The answer contains fundamental legal errors that completely misrepresent the law, legal principle, or legal concept. The answer would mislead a student and is factually wrong at its core.",
    "Irrelevant / NA": "The answer explicitly indicates that the information is unknown, unavailable, or not relevant to the question. This is appropriate when the AI correctly identifies that it cannot provide an answer."
}

# Captions for radio buttons (corresponding to ASSESSMENT_OPTIONS order)
ASSESSMENT_CAPTIONS = [
    SCORE_EXPLANATIONS["Perfect"],
    SCORE_EXPLANATIONS["Mostly correct"],
    SCORE_EXPLANATIONS["Noticeably flawed"],
    SCORE_EXPLANATIONS["Seriously wrong"],
    SCORE_EXPLANATIONS["Irrelevant / NA"]
]

def format_snake_case(text):
    """Convert snake_case to Title Case"""
    return ' '.join(word.capitalize() for word in text.split('_'))

def inject_tooltip_css():
    """Inject CSS to style radio button captions"""
    caption_css = """
    <style>
    /* Style captions to be subtle and informative */
    /* Target Streamlit's caption rendering - captions appear after radio button labels */
    .stRadio [data-testid="stMarkdownContainer"] p,
    .stRadio [data-testid="stMarkdownContainer"] span,
    .stRadio [data-testid="stMarkdownContainer"],
    .stRadio div[class*="caption"],
    .stRadio small,
    /* Target elements that come after radio button labels (where captions are rendered) */
    .stRadio > div > div[style*="margin"],
    .stRadio label + div,
    .stRadio label ~ div {{
        font-size: 0.85rem !important;
        color: #666 !important;
        line-height: 1.4 !important;
        margin-top: 2px !important;
        font-style: italic;
    }}
    
    /* More specific targeting for caption text */
    .stRadio [data-testid="stMarkdownContainer"] p {{
        margin-bottom: 8px !important;
        margin-top: 2px !important;
    }}
    </style>
    """
    st.markdown(caption_css, unsafe_allow_html=True)

def hash_password(password):
    """Hash a password using SHA256"""
    return hashlib.sha256(password.encode()).hexdigest()

@st.cache_data
def load_users(jurisdiction):
    """Load user credentials from Hugging Face Dataset for the specified jurisdiction"""
    try:
        file_path = hf_hub_download(
            repo_id=HF_DATASET_REPO,
            filename=f"{jurisdiction}/users/users.json",
            repo_type="dataset",
            token=HF_TOKEN
        )
        with open(file_path, 'r') as f:
            return json.load(f)
    except Exception:
        # File doesn't exist yet (first run), return empty dict
        return {}

def save_users(users, jurisdiction):
    """Save user credentials to Hugging Face Dataset for the specified jurisdiction"""
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
            json.dump(users, f, indent=2)
            temp_path = f.name
        
        hf_api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=f"{jurisdiction}/users/users.json",
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )
        os.unlink(temp_path)
        
        # Clear cache for users to ensure fresh data on next load
        load_users.clear(jurisdiction)
        
        return True
    except Exception as e:
        st.error(f"❌ **Error saving users to Hugging Face Dataset**: {str(e)}")
        raise

@st.cache_data(ttl=3600)  # Cache for 1 hour as safety measure
def load_user_data(username, jurisdiction):
    """Load user's answer data from Hugging Face Dataset for the specified jurisdiction"""
    try:
        file_path = hf_hub_download(
            repo_id=HF_DATASET_REPO,
            filename=f"{jurisdiction}/users/{username}_answers.parquet",
            repo_type="dataset",
            token=HF_TOKEN
        )
        return pd.read_parquet(file_path)
    except Exception:
        # File doesn't exist yet (new user), create new dataframe from grading template
        df = load_grading_template(jurisdiction)
        user_df = df.copy()
        user_df['legal_accuracy_score'] = None
        user_df['time_stamp'] = None
        return user_df

def save_user_data(username, user_df, jurisdiction, commit_message=None):
    """Save user's answer data to Hugging Face Dataset for the specified jurisdiction"""
    try:
        with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as f:
            user_df.to_parquet(f.name, index=False)
            temp_path = f.name
        
        upload_kwargs = {
            'path_or_fileobj': temp_path,
            'path_in_repo': f"{jurisdiction}/users/{username}_answers.parquet",
            'repo_id': HF_DATASET_REPO,
            'repo_type': "dataset",
            'token': HF_TOKEN
        }
        
        # Add commit_message if provided
        if commit_message:
            upload_kwargs['commit_message'] = commit_message
        
        hf_api.upload_file(**upload_kwargs)
        os.unlink(temp_path)
        
        # Clear cache for this user/jurisdiction to ensure fresh data on next load
        load_user_data.clear(username, jurisdiction)
        
        return True
    except Exception as e:
        st.error(f"❌ **Error saving user data to Hugging Face Dataset**: {str(e)}")
        raise

def update_user_answer(username, term, category, subcategory, question, answer, score, jurisdiction, df):
    """Update a specific answer in the user's data (deprecated - use update_category_answers for bulk updates)"""
    try:
        user_df = load_user_data(username, jurisdiction)
        
        # Find the matching row
        mask = (
            (user_df['term'] == term) &
            (user_df['category'] == category) &
            (user_df['subcategory'] == subcategory) &
            (user_df['question'] == question) &
            (user_df['answer'] == answer)
        )
        
        if mask.any():
            user_df.loc[mask, 'legal_accuracy_score'] = score
            save_user_data(username, user_df, jurisdiction)
            return True
        else:
            print(f"Warning: Could not find matching row for: {term}, {category}, {subcategory}, {question}")
            return False
    except Exception as e:
        print(f"Error updating answer: {str(e)}")
        return False

def auto_score_unknown_answers(username, term, category, df):
    """
    Automatically score all Unknown answers for a category.
    Returns list of (subcategory, question, answer, score) tuples.
    
    Args:
        username: User identifier
        term: Term name
        category: Category name
        df: Base dataframe with all questions/answers
    
    Returns:
        List of tuples (subcategory, question, answer, score) for Unknown answers
    """
    # Find all Unknown answers for this term-category pair (both "Unknown." and "Unknown")
    unknown_rows = df[(df['term'] == term) & 
                      (df['category'] == category) & 
                      ((df['answer'] == "Unknown.") | (df['answer'] == "Unknown"))]
    
    # Return list of tuples for bulk update (score is "NA" for Irrelevant / NA)
    return [(row['subcategory'], row['question'], row['answer'], "NA") 
            for _, row in unknown_rows.iterrows()]

def auto_score_all_unknown_answers_for_new_user(username, jurisdiction, df):
    """
    Automatically score all Unknown answers for all categories when a new user is created.
    This runs in the background during account creation.
    Performs a single bulk commit for all Unknown answers.
    """
    try:
        # Get all unique term-category pairs
        term_category_pairs = df[['term', 'category']].drop_duplicates().sort_values(['term', 'category']).values.tolist()
        
        # Collect all Unknown answers from all categories
        all_unknown_answers = []
        for term, category in term_category_pairs:
            unknown_answers = auto_score_unknown_answers(username, term, category, df)
            if unknown_answers:
                # Add term and category to each tuple for bulk update
                for subcategory, question, answer, score in unknown_answers:
                    all_unknown_answers.append((term, category, subcategory, question, answer, score))
        
        # If no Unknown answers found, return early
        if not all_unknown_answers:
            return True
        
        # Load user dataframe once
        user_df = load_user_data(username, jurisdiction)
        
        # Get current timestamp once for all updates
        current_timestamp = pd.Timestamp.now()
        
        # Update all Unknown answers in memory
        updated_count = 0
        for term, category, subcategory, question, answer, score in all_unknown_answers:
            # Find the matching row
            mask = (
                (user_df['term'] == term) &
                (user_df['category'] == category) &
                (user_df['subcategory'] == subcategory) &
                (user_df['question'] == question) &
                (user_df['answer'] == answer)
            )
            
            if mask.any():
                user_df.loc[mask, 'legal_accuracy_score'] = score
                user_df.loc[mask, 'time_stamp'] = current_timestamp
                updated_count += 1
            else:
                print(f"Warning: Could not find matching row for: {term}, {category}, {subcategory}, {question}")
        
        if updated_count == 0:
            print(f"Warning: No Unknown answers were updated for new user {username}")
            return False
        
        # Save once with a single commit message
        commit_message = f"Auto-score all Unknown answers for new user {username}"
        save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
        
        return True
    except Exception as e:
        print(f"Error auto-scoring Unknown answers for new user {username}: {str(e)}")
        return False

def update_category_answers(username, term, category, answers_list, jurisdiction, commit_message=None):
    """
    Update all answers for a category in a single commit.
    
    Args:
        username: User identifier
        term: Term name
        category: Category name
        answers_list: List of tuples (subcategory, question, answer, score)
        jurisdiction: Jurisdiction identifier
        commit_message: Optional commit message (auto-generated if None)
    
    Returns:
        True if successful, False otherwise
    """
    try:
        # Load user dataframe once
        user_df = load_user_data(username, jurisdiction)
        
        # Get current timestamp once for all updates in this category
        current_timestamp = pd.Timestamp.now()
        
        # Update all rows in memory
        updated_count = 0
        for subcategory, question, answer, score in answers_list:
            # Find the matching row
            mask = (
                (user_df['term'] == term) &
                (user_df['category'] == category) &
                (user_df['subcategory'] == subcategory) &
                (user_df['question'] == question) &
                (user_df['answer'] == answer)
            )
            
            if mask.any():
                user_df.loc[mask, 'legal_accuracy_score'] = score
                user_df.loc[mask, 'time_stamp'] = current_timestamp
                updated_count += 1
            else:
                print(f"Warning: Could not find matching row for: {term}, {category}, {subcategory}, {question}")
        
        if updated_count == 0:
            print(f"Warning: No rows were updated for {username} - {term} - {category}")
            return False
        
        # Generate commit message if not provided
        if commit_message is None:
            commit_message = f"Update answers for {username} - {term} - {category}"
        
        # Save once with commit message
        save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
        return True
    except Exception as e:
        print(f"Error updating category answers: {str(e)}")
        return False

def get_user_answer(username, term, category, subcategory, question, answer, jurisdiction):
    """Get user's answer for a specific question (checks both saved and pending answers)"""
    # First check if there's a pending answer in session state
    answer_key = (term, category, subcategory, question, answer)
    if answer_key in st.session_state.pending_term_answers:
        return st.session_state.pending_term_answers[answer_key]
    
    # Otherwise check saved data
    user_df = load_user_data(username, jurisdiction)
    
    mask = (
        (user_df['term'] == term) &
        (user_df['category'] == category) &
        (user_df['subcategory'] == subcategory) &
        (user_df['question'] == question) &
        (user_df['answer'] == answer)
    )
    
    if mask.any():
        score = user_df.loc[mask, 'legal_accuracy_score'].iloc[0]
        if pd.notna(score):
            return score
    return None

def find_first_unanswered_category(username, jurisdiction, df):
    """Find the first category that hasn't been fully answered"""
    user_df = load_user_data(username, jurisdiction)
    
    # Get term_category_pairs for this jurisdiction
    term_category_pairs = get_term_category_pairs(df)
    
    for idx, (term, category) in enumerate(term_category_pairs):
        # Get all subcategories for this term-category pair from base df
        category_data = df[(df['term'] == term) & (df['category'] == category)]
        
        # Check if all questions in this category are answered
        all_answered = True
        for _, row in category_data.iterrows():
            # Check directly in user_df instead of calling get_user_answer (which reloads data)
            mask = (
                (user_df['term'] == row['term']) &
                (user_df['category'] == row['category']) &
                (user_df['subcategory'] == row['subcategory']) &
                (user_df['question'] == row['question']) &
                (user_df['answer'] == row['answer'])
            )
            
            if mask.any():
                score = user_df.loc[mask, 'legal_accuracy_score'].iloc[0]
                if pd.isna(score):
                    all_answered = False
                    break
            else:
                # No matching row found - question not answered
                all_answered = False
                break
        
        if not all_answered:
            return idx
    
    return len(term_category_pairs)  # All answered, return last index

def restore_submitted_status(username, jurisdiction, df):
    """Restore submitted status for categories that have all answers in parquet file"""
    user_df = load_user_data(username, jurisdiction)
    
    # Get term_category_pairs for this jurisdiction
    term_category_pairs = get_term_category_pairs(df)
    
    submitted_pairs = set()
    for idx, (term, category) in enumerate(term_category_pairs):
        pair_key = f"{term}_{category}_{idx}"
        
        # Get all subcategories for this term-category pair
        category_data = df[(df['term'] == term) & (df['category'] == category)]
        
        # Check if all questions in this category are answered
        all_answered = True
        for _, row in category_data.iterrows():
            # Check directly in user_df instead of calling get_user_answer (which reloads data)
            mask = (
                (user_df['term'] == row['term']) &
                (user_df['category'] == row['category']) &
                (user_df['subcategory'] == row['subcategory']) &
                (user_df['question'] == row['question']) &
                (user_df['answer'] == row['answer'])
            )
            
            if mask.any():
                score = user_df.loc[mask, 'legal_accuracy_score'].iloc[0]
                if pd.isna(score):
                    all_answered = False
                    break
            else:
                # No matching row found - question not answered
                all_answered = False
                break
        
        if all_answered:
            submitted_pairs.add(pair_key)
    
    return submitted_pairs

def category_has_subcategories(term, category_name, df):
    """
    Check if a category has any subcategories after filtering Unknown answers.
    
    Args:
        term: Term name
        category_name: Category name
        df: Base dataframe with all questions/answers
    
    Returns:
        True if category has at least one subcategory after filtering Unknown answers, False otherwise
    """
    # Filter out subcategories where answer is "Unknown." or "Unknown"
    filtered_df = df[(df['term'] == term) & 
                    (df['category'] == category_name) & 
                    (df['answer'] != "Unknown.") & 
                    (df['answer'] != "Unknown")]
    
    # Check if there are any subcategories after filtering
    subcategory_names = filtered_df['subcategory'].unique()
    return len(subcategory_names) > 0

class Subcategory:
    """Represents a single subcategory with its question, answer, and radio button"""
    
    def __init__(self, term, category, subcategory_name, df):
        self.term = term
        self.category = category
        self.subcategory_name = subcategory_name
        self.formatted_name = format_snake_case(subcategory_name)
        
        # Get question and answer
        result = df[(df['term'] == term) & 
                   (df['category'] == category) & 
                   (df['subcategory'] == subcategory_name)]
        if len(result) > 0:
            self.question = result.iloc[0]['question']
            self.answer = result.iloc[0]['answer']
        else:
            self.question = None
            self.answer = None


class Category:
    """Represents a category with its subcategories"""
    
    def __init__(self, term, category_name, df):
        self.term = term
        self.category_name = category_name
        self.formatted_name = format_snake_case(category_name)
        
        # Filter out subcategories where answer is "Unknown." or "Unknown"
        filtered_df = df[(df['term'] == term) & 
                        (df['category'] == category_name) & 
                        (df['answer'] != "Unknown.") & 
                        (df['answer'] != "Unknown")]
        
        # Get all subcategories for this term-category pair (excluding Unknown answers)
        # Sort by subcategory_index to maintain the configured order
        subcat_data = filtered_df[['subcategory', 'subcategory_index']].drop_duplicates()
        subcat_data = subcat_data.sort_values('subcategory_index')
        subcategory_names = subcat_data['subcategory'].tolist()
        
        # Create Subcategory instances (only for non-Unknown answers)
        self.subcategories = [
            Subcategory(term, category_name, subcat_name, df)
            for subcat_name in subcategory_names
        ]


class Term:
    """Represents a term with its categories"""
    
    def __init__(self, term_name, df):
        self.term_name = term_name
        self.formatted_name = format_snake_case(term_name)
        
        # Get all categories for this term, sorted by category_index
        cat_data = df[df['term'] == term_name][['category', 'category_index']].drop_duplicates()
        cat_data = cat_data.sort_values('category_index')
        category_names = cat_data['category'].tolist()
        
        # Create Category instances
        self.categories = [
            Category(term_name, cat_name, df)
            for cat_name in category_names
        ]
    
    def get_category_by_name(self, category_name):
        """Get a category by its name"""
        for cat in self.categories:
            if cat.category_name == category_name:
                return cat
        return None


@st.cache_data
def get_term_category_pairs(df):
    """Get filtered term-category pairs, cached to avoid recomputation on every rerun"""
    # Get all unique term-category pairs with their category indexes
    all_pairs_df = df[['term', 'category', 'category_index']].drop_duplicates()
    
    # Sort by term name and category_index
    all_pairs_df = all_pairs_df.sort_values(['term', 'category_index'])
    
    # Convert to list efficiently (avoids slow iterrows)
    all_pairs_list = all_pairs_df[['term', 'category']].values.tolist()
    
    # Filter out categories that have no subcategories after filtering Unknown answers
    filtered_pairs = [(term, category) for term, category in all_pairs_list 
                      if category_has_subcategories(term, category, df)]
    
    return filtered_pairs

# Cache for Term instances (keyed by jurisdiction and term_name)
term_cache = {}

def get_term_instance(term_name, df):
    """Get or create a Term instance for the given dataframe"""
    cache_key = f"{id(df)}_{term_name}"  # Use df id to differentiate jurisdictions
    if cache_key not in term_cache:
        term_cache[cache_key] = Term(term_name, df)
    return term_cache[cache_key]

def get_category_for_pair(term_name, category_name, df):
    """Get Category instance for a term-category pair"""
    term = get_term_instance(term_name, df)
    return term.get_category_by_name(category_name)

# Initialize session state
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False
if 'username' not in st.session_state:
    st.session_state.username = None
if 'jurisdiction' not in st.session_state:
    st.session_state.jurisdiction = None
if 'current_index' not in st.session_state:
    st.session_state.current_index = 0
if 'show_term_complete' not in st.session_state:
    st.session_state.show_term_complete = False
if 'completed_term' not in st.session_state:
    st.session_state.completed_term = None
if 'next_term' not in st.session_state:
    st.session_state.next_term = None
if 'show_term_back_warning' not in st.session_state:
    st.session_state.show_term_back_warning = False
if 'back_current_term' not in st.session_state:
    st.session_state.back_current_term = None
if 'back_previous_term' not in st.session_state:
    st.session_state.back_previous_term = None
if 'back_current_index' not in st.session_state:
    st.session_state.back_current_index = None
if 'show_guide' not in st.session_state:
    st.session_state.show_guide = True
if 'submitted_pairs' not in st.session_state:
    st.session_state.submitted_pairs = set()  # Track which pairs have been submitted
if 'original_selections' not in st.session_state:
    st.session_state.original_selections = {}  # Store original selections for each pair
if 'has_unsaved_changes' not in st.session_state:
    st.session_state.has_unsaved_changes = {}  # Track unsaved changes per pair
if 'pending_term_answers' not in st.session_state:
    st.session_state.pending_term_answers = {}  # Store uncommitted answers: {(term, category, subcategory, question, answer): score}
if 'current_term_name' not in st.session_state:
    st.session_state.current_term_name = None  # Track which term is being worked on

# Login page
if not st.session_state.logged_in:
    st.markdown("# Login")
    st.markdown("Please select a jurisdiction and enter your username and password to continue.")
    
    # Jurisdiction selector - discover available jurisdictions dynamically
    available_jurisdictions = discover_available_jurisdictions()
    if not available_jurisdictions:
        st.error("❌ **Error**: No jurisdictions found in the repository. Please ensure the repository structure is correct.")
        st.stop()
    
    # Get display name mapping
    display_mapping = get_jurisdiction_display_mapping(available_jurisdictions)
    
    # Determine default index for selectbox - prioritize hr-hr
    default_index = 0
    if "hr-hr" in available_jurisdictions:
        default_index = available_jurisdictions.index("hr-hr")
        # Set hr-hr as default in session state if not already set
        if not st.session_state.jurisdiction:
            st.session_state.jurisdiction = "hr-hr"
    elif st.session_state.jurisdiction and st.session_state.jurisdiction in available_jurisdictions:
        default_index = available_jurisdictions.index(st.session_state.jurisdiction)
    
    # Create format function to show display names
    def format_jurisdiction(code):
        return display_mapping.get(code, code)
    
    jurisdiction = st.selectbox(
        "Jurisdiction", 
        options=available_jurisdictions, 
        index=default_index,
        format_func=format_jurisdiction
    )
    st.session_state.jurisdiction = jurisdiction
    
    username = st.text_input("Username")
    password = st.text_input("Password", type="password")
    
    col1, col2 = st.columns(2)
    with col1:
        if st.button("Login", type="primary", use_container_width=True):
            if not jurisdiction:
                st.error("Please select a jurisdiction")
            else:
                users = load_users(jurisdiction)
                
                if username in users:
                    # Existing user - check password
                    if users[username]['password'] == hash_password(password):
                        st.session_state.logged_in = True
                        st.session_state.username = username
                        # Load grading template for this jurisdiction
                        df = load_grading_template(jurisdiction)
                        # Restore submitted status for previously submitted categories
                        st.session_state.submitted_pairs = restore_submitted_status(username, jurisdiction, df)
                        # Find first unanswered category and resume there
                        resume_index = find_first_unanswered_category(username, jurisdiction, df)
                        st.session_state.current_index = resume_index
                        st.rerun()
                    else:
                        st.error("Incorrect password")
                else:
                    # Username not found - require registration
                    st.error("Username not found. Please register first using the 'Register New User' button.")
    
    with col2:
        if st.button("Register New User", use_container_width=True):
            if not jurisdiction:
                st.error("Please select a jurisdiction")
            else:
                users = load_users(jurisdiction)
                if username in users:
                    st.error("Username already exists")
                elif username and password:
                    users[username] = {'password': hash_password(password)}
                    save_users(users, jurisdiction)
                    # Load grading template for this jurisdiction
                    df = load_grading_template(jurisdiction)
                    # Auto-score all Unknown answers for the new user in the background
                    auto_score_all_unknown_answers_for_new_user(username, jurisdiction, df)
                    st.success("User registered successfully! Please click Login.")
                else:
                    st.error("Please enter both username and password")

# Main application (only shown if logged in)
elif st.session_state.logged_in:
    username = st.session_state.username
    jurisdiction = st.session_state.jurisdiction
    current_index = st.session_state.current_index
    
    # Safety check: ensure jurisdiction is set (handles case where user was logged in before jurisdiction feature was added)
    available_jurisdictions = discover_available_jurisdictions()
    if not jurisdiction or jurisdiction not in available_jurisdictions:
        st.error("❌ **Error**: Jurisdiction is not set. Please log out and log in again with a jurisdiction selected.")
        if st.button("Logout"):
            st.session_state.logged_in = False
            st.session_state.username = None
            st.session_state.jurisdiction = None
            st.session_state.current_index = 0
            st.session_state.show_guide = True
            st.session_state.submitted_pairs = set()
            st.session_state.original_selections = {}
            st.session_state.has_unsaved_changes = {}
            st.rerun()
        st.stop()
    
    # Load grading template for the selected jurisdiction
    df = load_grading_template(jurisdiction)
    
    # Get term_category_pairs for this jurisdiction
    term_category_pairs = get_term_category_pairs(df)
    total_pairs = len(term_category_pairs)
    
    # Debug info (can be removed in production)
    with st.sidebar:
        with st.expander("Debug Info"):
            st.write(f"HF Dataset Repo: `{HF_DATASET_REPO}`")
            st.write(f"HF Token configured: {HF_TOKEN is not None}")
            st.write(f"HF API initialized: {hf_api is not None}")
            if username:
                jurisdiction_display = get_jurisdiction_display_name(jurisdiction)
                st.write(f"Jurisdiction: {jurisdiction_display} (`{jurisdiction}`)")
                st.write(f"User parquet file: `{jurisdiction}/users/{username}_answers.parquet`")
                st.write(f"Users file: `{jurisdiction}/users/users.json`")
    
    # Check if we should show the annotation guide first
    if st.session_state.show_guide:
        # Annotation Guide Page
        st.markdown("# Welcome!")
        st.markdown("")
        st.markdown("## Task Overview")
        st.markdown("""
        You are a legal expert tasked with evaluating AI-generated structured answers to 55 specific legal questions from court opinions. 
        Your task is to assess the legal accuracy for each question and ensure they meet the high standards required for legal education.
        
        For each question, you will be presented with:
        - **The legal question** that was asked
        - **The AI-generated answer** that needs to be evaluated
        
        Your role is to critically evaluate whether the answer is legally accurate based on your knowledge of the law. Consider:
        - Does the answer correctly state the legal principle or rule?
        - Are there any factual inaccuracies or misstatements of law?
        - Is the answer complete and does it address the question properly?
        - Would this answer be acceptable for use in legal education?
        """)
        st.markdown("")
        st.markdown("## Grading Criteria")
        st.markdown("""
        Based on your general knowledge of the law, consider how legally accurate the answer is and grade it as follows:
        
        - **Perfect** - Only when truly flawless. The answer is legally accurate, well-stated, and appropriate for legal education. It correctly explains the legal principle, rule, or concept without any discernible errors or misleading statements.
        
        - **Mostly correct** - One very small issue or slightly awkward phrasing. The answer is generally accurate but contains minor inaccuracies, imprecise language, or could be more precise. The core legal content is correct, but there are small issues that could be improved for educational purposes.
        
        - **Noticeably flawed** - Clear error but main idea still ok. The answer contains significant errors that substantially affect its accuracy. While not completely wrong, there are important mistakes in legal reasoning, application of law, or factual statements that would confuse or mislead students.
        
        - **Seriously wrong** - Hallucination, wrong format, meaningless, etc. The answer contains fundamental legal errors that completely misrepresent the law, legal principle, or legal concept. The answer would mislead a student and is factually wrong at its core.
        
        - **Irrelevant / NA** - The answer explicitly indicates that the information is unknown, unavailable, or not relevant to the question. This is appropriate when the AI correctly identifies that it cannot provide an answer.
        """)
        st.markdown("")
        st.markdown("## Examples")
        st.markdown("""
        ### Example 1: Contract Formation
        
        **Question:** What is required for a valid contract to be formed?
        
        **AI Answer:** "A valid contract requires an offer, acceptance, consideration, and mutual assent between the parties. The parties must have the legal capacity to contract, and the subject matter must be legal."
        
        **Evaluation:** This answer is **Perfect**. It correctly identifies the essential elements of contract formation (offer, acceptance, consideration) and mentions additional important requirements (capacity, legality). The answer is comprehensive and legally accurate.
        
        ---
        
        **AI Answer (Alternative):** "A contract is formed when two people agree to something and shake hands on it."
        
        **Evaluation:** This answer would be **Seriously wrong**. While it mentions agreement, it completely misrepresents contract law by suggesting that a handshake alone is sufficient. It omits essential elements like consideration and fails to explain the legal requirements properly.
        
        ### Example 2: Due Process Rights
        
        **Question:** What procedural protections are required under the Due Process Clause of the Fourteenth Amendment?
        
        **AI Answer:** "The Due Process Clause requires that individuals receive notice of proceedings against them and an opportunity to be heard. This includes the right to present evidence, cross-examine witnesses, and have a neutral decision-maker. The specific procedures required depend on the nature of the interest at stake and the context of the deprivation."
        
        **Evaluation:** This answer is **Perfect**. It correctly identifies the core elements of procedural due process (notice and opportunity to be heard) and mentions the balancing test that courts apply. The answer is legally sound.
        
        ---
        
        **AI Answer (Alternative):** "Due process means you get a fair trial with a lawyer, a jury, and the right to remain silent, just like in criminal cases."
        
        **Evaluation:** This answer would be **Noticeably flawed**. While it mentions some procedural protections, it incorrectly conflates criminal procedure rights (jury trial, right to remain silent) with the broader concept of due process, which applies to civil proceedings as well. The answer oversimplifies and contains significant inaccuracies about when these specific rights apply.
        """)
        st.markdown("")
        st.markdown("### You've got this! Your expertise is invaluable for improving legal education.")
        st.markdown("")  # Spacing
        col1, col2, col3 = st.columns([1, 1, 1])
        with col2:
            if st.button("Start Assessment", type="primary", use_container_width=True):
                st.session_state.show_guide = False
                st.rerun()
    
    # Check if we should show the back warning page (moving to different term)
    elif st.session_state.show_term_back_warning:
        # Show warning page when going back to a different term
        st.markdown(
            f"""
            <div style='text-align: center; font-size: 1.2em;'>
                <p>You are leaving the current term</p>
                <p><strong>{st.session_state.back_current_term}</strong></p>
                <br>
                <p>and going back to questions about the term</p>
                <p><strong>{st.session_state.back_previous_term}</strong></p>
            </div>
            """,
            unsafe_allow_html=True
        )
        st.markdown("")  # Spacing
        col1, col2, col3 = st.columns([1, 1, 1])
        with col1:
            if st.button("Continue to Previous Term", type="primary", use_container_width=True):
                # Move to the previous term-category pair
                st.session_state.current_index -= 1
                st.session_state.show_term_back_warning = False
                st.session_state.back_current_term = None
                st.session_state.back_previous_term = None
                st.session_state.back_current_index = None
                st.rerun()
        with col3:
            if st.button(f"Stay on {st.session_state.back_current_term}", use_container_width=True):
                # Cancel navigation - stay on current term
                st.session_state.current_index = st.session_state.back_current_index
                st.session_state.show_term_back_warning = False
                st.session_state.back_current_term = None
                st.session_state.back_previous_term = None
                st.session_state.back_current_index = None
                st.rerun()
    
    # Check if we should show the term completion page (moving forward to different term)
    elif st.session_state.show_term_complete:
        # Show intermediate page between terms - centered with terms on new lines and bold
        st.markdown(
            f"""
            <div style='text-align: center; font-size: 1.2em;'>
                <p>You finished the assessment for all questions for the term</p>
                <p><strong>{st.session_state.completed_term}</strong></p>
                <br>
                <p>The next term to be assessed is</p>
                <p><strong>{st.session_state.next_term}</strong></p>
            </div>
            """,
            unsafe_allow_html=True
        )
        st.markdown("")  # Spacing
        col1, col2, col3 = st.columns([1, 1, 1])
        with col2:
            if st.button("Continue to Next Term", type="primary", use_container_width=True):
                # Move to the next term-category pair
                st.session_state.current_index += 1
                st.session_state.show_term_complete = False
                st.session_state.completed_term = None
                st.session_state.next_term = None
                st.rerun()
    
    elif current_index < total_pairs:
        term_name, category_name = term_category_pairs[current_index]
        category = get_category_for_pair(term_name, category_name, df)
        term = get_term_instance(term_name, df)
        
        # Safety check: skip if category has no subcategories (shouldn't happen due to filtering, but just in case)
        if not category or len(category.subcategories) == 0:
            # Skip to next category
            st.session_state.current_index += 1
            st.rerun()
        
        # Display header
        st.markdown(f"# Term: {term.formatted_name}")
        st.markdown(f"## Category: {category.formatted_name}")
        
        # Show indicator if there are pending uncommitted answers for this term
        if st.session_state.pending_term_answers:
            pending_count = sum(1 for key in st.session_state.pending_term_answers.keys() if key[0] == term_name)
            if pending_count > 0:
                st.info(f"💾 {pending_count} answer{'s' if pending_count > 1 else ''} pending (will be saved when term is complete)")
        
        # Display all subcategories with their questions, answers, and radio buttons
        pair_key = f"{term_name}_{category_name}_{current_index}"
        is_submitted = pair_key in st.session_state.submitted_pairs
        
        # Don't pre-load values into session_state - let the radio buttons manage their own state
        # This prevents the warning about default value vs Session State API conflict
        
        # Check if category is fully answered in parquet file
        # This includes both visible subcategories and Unknown answers (which are auto-scored)
        category_fully_answered = True
        
        # Check visible subcategories
        for i, subcat in enumerate(category.subcategories):
            saved_score = get_user_answer(username, term_name, category_name, subcat.subcategory_name,
                                         subcat.question, subcat.answer, jurisdiction)
            if saved_score is None:
                category_fully_answered = False
                break
        
        # Also check Unknown answers (they should be auto-scored as "NA")
        if category_fully_answered:
            unknown_answers = auto_score_unknown_answers(username, term_name, category_name, df)
            for subcategory, question, answer, score in unknown_answers:
                saved_score = get_user_answer(username, term_name, category_name, subcategory, question, answer, jurisdiction)
                if saved_score is None:
                    category_fully_answered = False
                    break
        
        # Update is_submitted based on parquet file check
        if category_fully_answered and not is_submitted:
            # Category is fully answered in parquet but not marked as submitted in session
            # Mark it as submitted for UI purposes (to show Update button)
            is_submitted = True
            st.session_state.submitted_pairs.add(pair_key)
        
        # Load original selections if this pair has been submitted
        if is_submitted:
            if pair_key not in st.session_state.original_selections:
                st.session_state.original_selections[pair_key] = {}
                for i, subcat in enumerate(category.subcategories):
                    radio_key = f"{pair_key}_subcat_{i}"
                    # Get from session_state (which was just set above) or from saved data
                    if radio_key in st.session_state:
                        st.session_state.original_selections[pair_key][radio_key] = st.session_state[radio_key]
                    else:
                        saved_score = get_user_answer(username, term_name, category_name, subcat.subcategory_name,
                                                     subcat.question, subcat.answer, jurisdiction)
                        if saved_score is not None:
                            score_to_option = {v: k for k, v in ASSESSMENT_TO_SCORE.items()}
                            if saved_score in score_to_option:
                                st.session_state.original_selections[pair_key][radio_key] = score_to_option[saved_score]
        
        all_selected = True
        changed_count = 0
        changes_info = []  # Store info about changes for display
        selected_values = {}  # Store all selected values for validation
        
        # Inject tooltip CSS and JavaScript for score explanations
        inject_tooltip_css()
        
        for i, subcat in enumerate(category.subcategories):
            st.markdown(f"**Subcategory: {subcat.formatted_name}**")
            
            # Display question and answer in a styled box
            st.markdown(
                f"""
                <div style='border: 2px solid #ccc; padding: 15px; margin: 10px 0; border-radius: 5px; background-color: #f9f9f9; color: #000;'>
                    <strong style='color: #000;'>Question:</strong> <span style='color: #000;'>{subcat.question}</span><br><br>
                    <strong style='color: #000;'>Answer:</strong> <span style='color: #000;'>{subcat.answer}</span>
                </div>
                """,
                unsafe_allow_html=True
            )
            
            # Selectbox for this subcategory
            radio_key = f"{pair_key}_subcat_{i}"
            
            # Get original value for change detection
            original_value = None
            if is_submitted and pair_key in st.session_state.original_selections:
                original_value = st.session_state.original_selections[pair_key].get(radio_key)
            
            # Get saved value to determine if we should set a default index
            # Don't pre-set session_state - only use index parameter to avoid conflicts
            saved_score = get_user_answer(username, term_name, category_name, subcat.subcategory_name,
                                         subcat.question, subcat.answer, jurisdiction)
            default_index = None
            
            # Check if there's a saved value in parquet file
            if saved_score is not None:
                score_to_option = {v: k for k, v in ASSESSMENT_TO_SCORE.items()}
                if saved_score in score_to_option:
                    saved_option = score_to_option[saved_score]
                    if saved_option in ASSESSMENT_OPTIONS:
                        default_index = ASSESSMENT_OPTIONS.index(saved_option)
            # If no saved value but value exists in session_state (from user interaction), use it
            elif radio_key in st.session_state:
                current_value = st.session_state[radio_key]
                if current_value in ASSESSMENT_OPTIONS:
                    default_index = ASSESSMENT_OPTIONS.index(current_value)
            
            # Use radio buttons with index=None when no saved value to prevent default selection
            # This prevents the warning about default value vs Session State API
            if default_index is not None:
                # Has saved value - set index to that value
                selected_value = st.radio(
                    "Your Assessment:",
                    options=ASSESSMENT_OPTIONS,
                    captions=ASSESSMENT_CAPTIONS,
                    key=radio_key,
                    index=default_index
                )
            else:
                # No saved value - use index=None to prevent any default selection
                # Radio button will return None until user makes a selection
                selected_value = st.radio(
                    "Your Assessment:",
                    options=ASSESSMENT_OPTIONS,
                    captions=ASSESSMENT_CAPTIONS,
                    key=radio_key,
                    index=None
                )
            
            # Store the selected value for later validation
            selected_values[radio_key] = selected_value
            
            # Check for changes if submitted
            if is_submitted and original_value is not None:
                if selected_value != original_value:
                    changed_count += 1
                    changes_info.append({
                        'subcategory': subcat.formatted_name,
                        'old': original_value,
                        'new': selected_value
                    })
                    # Display change indicator
                    if selected_value is None:
                        st.error(f"⚠️ **Invalid selection**: Please select a valid assessment.")
                    elif selected_value is not None:
                        st.info(f"⚠️ Changed from: **{original_value}** → **{selected_value}**")
            
            # Add separator between subcategories (except for the last one)
            if i < len(category.subcategories) - 1:
                st.markdown("---")
        
        # Track unsaved changes
        st.session_state.has_unsaved_changes[pair_key] = changed_count > 0
        
        # Navigation buttons
        st.markdown("")  # Spacing
        col1, col2, col3 = st.columns([1, 1, 1])
        
        # Back button
        with col1:
            can_go_back = current_index > 0
            if st.button("← Back", disabled=not can_go_back, use_container_width=True):
                # Check for unsaved changes
                if st.session_state.has_unsaved_changes.get(pair_key, False):
                    st.warning("⚠️ You have unsaved changes. Please click 'Update' before navigating.")
                else:
                    # Check if moving to a different term
                    prev_index = current_index - 1
                    if prev_index >= 0:
                        prev_term_name, _ = term_category_pairs[prev_index]
                        if prev_term_name != term_name:
                            # Moving to a different term - show warning page
                            # Store current index so we can restore it if user cancels
                            st.session_state.back_current_index = current_index
                            st.session_state.show_term_back_warning = True
                            st.session_state.back_current_term = term.formatted_name
                            prev_term = get_term_instance(prev_term_name, df)
                            st.session_state.back_previous_term = prev_term.formatted_name
                        else:
                            # Same term, just move back
                            st.session_state.current_index = prev_index
                            st.session_state.show_term_complete = False
                            st.session_state.show_term_back_warning = False
                    else:
                        # Can't go back further
                        st.session_state.current_index = prev_index
                        st.session_state.show_term_complete = False
                        st.session_state.show_term_back_warning = False
                    st.rerun()
        
        # Submit/Update button
        with col2:
            # Check that all selections are valid (not None)
            # Use session_state as the source of truth since it's always up-to-date with current widget state
            all_valid = True
            for i, subcat in enumerate(category.subcategories):
                radio_key = f"{pair_key}_subcat_{i}"
                # Get current value from session_state (this reflects the actual current selection)
                selected_value = st.session_state.get(radio_key)
                # If value is None, it's not valid
                if selected_value is None:
                    all_valid = False
                    break
            
            if all_valid and len(category.subcategories) > 0:
                if is_submitted and changed_count > 0:
                    # Update button with count
                    button_label = f"Update ({changed_count} answer{'s' if changed_count > 1 else ''})"
                    if st.button(button_label, type="primary", use_container_width=True):
                        # Collect all answers and store in session state (don't commit yet)
                        for i, subcat in enumerate(category.subcategories):
                            radio_key = f"{pair_key}_subcat_{i}"
                            selected_value = st.session_state.get(radio_key)
                            if selected_value is not None:
                                score = ASSESSMENT_TO_SCORE[selected_value]
                                answer_key = (term_name, category_name, subcat.subcategory_name, subcat.question, subcat.answer)
                                st.session_state.pending_term_answers[answer_key] = score
                        
                        # Save current selections as new originals
                        st.session_state.original_selections[pair_key] = {}
                        for i in range(len(category.subcategories)):
                            radio_key = f"{pair_key}_subcat_{i}"
                            if radio_key in st.session_state:
                                st.session_state.original_selections[pair_key][radio_key] = st.session_state[radio_key]
                        st.session_state.has_unsaved_changes[pair_key] = False
                        st.success("Answers updated! (Will be saved when term is complete)")
                        st.rerun()
                elif is_submitted:
                    # Already submitted, no changes
                    st.button("Update (0 answers)", disabled=True, use_container_width=True)
                else:
                    # Submit button for first time
                    if st.button("Submit", type="primary", use_container_width=True):
                        # Collect all user answers and store in session state (don't commit yet)
                        for i, subcat in enumerate(category.subcategories):
                            radio_key = f"{pair_key}_subcat_{i}"
                            selected_value = st.session_state.get(radio_key)
                            if selected_value is not None:
                                score = ASSESSMENT_TO_SCORE[selected_value]
                                answer_key = (term_name, category_name, subcat.subcategory_name, subcat.question, subcat.answer)
                                st.session_state.pending_term_answers[answer_key] = score
                        
                        # Mark as submitted and save original selections
                        st.session_state.submitted_pairs.add(pair_key)
                        st.session_state.original_selections[pair_key] = {}
                        for i in range(len(category.subcategories)):
                            radio_key = f"{pair_key}_subcat_{i}"
                            if radio_key in st.session_state:
                                st.session_state.original_selections[pair_key][radio_key] = st.session_state[radio_key]
                        
                        # Set current term name
                        st.session_state.current_term_name = term_name
                        
                        # Check if this is the last category for the current term
                        current_term_name = term_name
                        next_index = current_index + 1
                        
                        if next_index < total_pairs:
                            next_term_name, _ = term_category_pairs[next_index]
                            
                            # Check if we're moving to a different term
                            if next_term_name != current_term_name:
                                # Commit all pending answers for this term before switching
                                if st.session_state.pending_term_answers:
                                    # Load user dataframe
                                    user_df = load_user_data(username, jurisdiction)
                                    current_timestamp = pd.Timestamp.now()
                                    
                                    # Update all pending answers
                                    for answer_key, score in st.session_state.pending_term_answers.items():
                                        t, c, sc, q, a = answer_key
                                        mask = (
                                            (user_df['term'] == t) &
                                            (user_df['category'] == c) &
                                            (user_df['subcategory'] == sc) &
                                            (user_df['question'] == q) &
                                            (user_df['answer'] == a)
                                        )
                                        if mask.any():
                                            user_df.loc[mask, 'legal_accuracy_score'] = score
                                            user_df.loc[mask, 'time_stamp'] = current_timestamp
                                    
                                    # Commit all changes for this term
                                    commit_message = f"Complete term {current_term_name} - {username}"
                                    save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
                                    
                                    # Clear pending answers
                                    st.session_state.pending_term_answers = {}
                                
                                # Show intermediate page
                                st.session_state.show_term_complete = True
                                st.session_state.completed_term = term.formatted_name
                                next_term = get_term_instance(next_term_name, df)
                                st.session_state.next_term = next_term.formatted_name
                            else:
                                # Same term, just move to next category
                                st.session_state.current_index = next_index
                        else:
                            # No more pairs - commit any pending answers
                            if st.session_state.pending_term_answers:
                                # Load user dataframe
                                user_df = load_user_data(username, jurisdiction)
                                current_timestamp = pd.Timestamp.now()
                                
                                # Update all pending answers
                                for answer_key, score in st.session_state.pending_term_answers.items():
                                    t, c, sc, q, a = answer_key
                                    mask = (
                                        (user_df['term'] == t) &
                                        (user_df['category'] == c) &
                                        (user_df['subcategory'] == sc) &
                                        (user_df['question'] == q) &
                                        (user_df['answer'] == a)
                                    )
                                    if mask.any():
                                        user_df.loc[mask, 'legal_accuracy_score'] = score
                                        user_df.loc[mask, 'time_stamp'] = current_timestamp
                                
                                # Commit all changes for this term
                                commit_message = f"Complete term {current_term_name} - {username}"
                                save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
                                
                                # Clear pending answers
                                st.session_state.pending_term_answers = {}
                            
                            # Go to finish
                            st.session_state.current_index = next_index
                        
                        st.success("Answers saved! (Will be committed when term is complete)")
                        st.rerun()
            else:
                st.button("Submit", disabled=True, use_container_width=True)
        
        # Forward button
        with col3:
            # Check if we would be moving to a different term
            next_index = current_index + 1
            moving_to_different_term = False
            term_is_complete = False
            
            if next_index < total_pairs:
                next_term_name, _ = term_category_pairs[next_index]
                moving_to_different_term = (next_term_name != term_name)
                
                if moving_to_different_term:
                    # Check if all categories in current term are complete
                    term_is_complete = True
                    for idx, (t, c) in enumerate(term_category_pairs):
                        if t == term_name:
                            pk = f"{t}_{c}_{idx}"
                            if pk not in st.session_state.submitted_pairs:
                                term_is_complete = False
                                break
            
            # Can only go forward if current pair is submitted and no unsaved changes
            # If moving to different term, also require that current term is complete
            can_go_forward = (pair_key in st.session_state.submitted_pairs and 
                             not st.session_state.has_unsaved_changes.get(pair_key, False) and
                             current_index < total_pairs - 1 and
                             (not moving_to_different_term or term_is_complete))
            
            forward_button_disabled = not can_go_forward
            
            if st.button("Forward →", disabled=forward_button_disabled, use_container_width=True):
                if st.session_state.has_unsaved_changes.get(pair_key, False):
                    st.warning("⚠️ You have unsaved changes. Please click 'Update' before navigating.")
                else:
                    # Check if moving to a different term
                    if next_index < total_pairs:
                        next_term_name, _ = term_category_pairs[next_index]
                        if next_term_name != term_name:
                            # Commit all pending answers for current term before switching
                            if st.session_state.pending_term_answers:
                                # Load user dataframe
                                user_df = load_user_data(username, jurisdiction)
                                current_timestamp = pd.Timestamp.now()
                                
                                # Update all pending answers
                                for answer_key, score in st.session_state.pending_term_answers.items():
                                    t, c, sc, q, a = answer_key
                                    mask = (
                                        (user_df['term'] == t) &
                                        (user_df['category'] == c) &
                                        (user_df['subcategory'] == sc) &
                                        (user_df['question'] == q) &
                                        (user_df['answer'] == a)
                                    )
                                    if mask.any():
                                        user_df.loc[mask, 'legal_accuracy_score'] = score
                                        user_df.loc[mask, 'time_stamp'] = current_timestamp
                                
                                # Commit all changes for this term
                                commit_message = f"Complete term {term_name} - {username}"
                                save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
                                
                                # Clear pending answers
                                st.session_state.pending_term_answers = {}
                            
                            # Moving to a different term - show term switching page
                            st.session_state.show_term_complete = True
                            st.session_state.completed_term = term.formatted_name
                            next_term = get_term_instance(next_term_name, df)
                            st.session_state.next_term = next_term.formatted_name
                        else:
                            # Same term, just move forward
                            st.session_state.current_index = next_index
                            st.session_state.show_term_complete = False
                    else:
                        # No more pairs
                        st.session_state.current_index = next_index
                        st.session_state.show_term_complete = False
                    st.session_state.show_term_back_warning = False
                    st.rerun()
            
            # Show helpful message if forward is blocked due to incomplete term
            if moving_to_different_term and not term_is_complete:
                st.caption("⚠️ Complete all categories in this term before proceeding")
    
    else:
        # Finished
        st.markdown("# You are finished.")
        st.markdown("## Thank you for your contribution!")
        
        # Logout button
        if st.button("Logout"):
            st.session_state.logged_in = False
            st.session_state.username = None
            st.session_state.jurisdiction = None
            st.session_state.current_index = 0
            st.session_state.show_guide = True
            st.session_state.submitted_pairs = set()
            st.session_state.original_selections = {}
            st.session_state.has_unsaved_changes = {}
            st.rerun()