"""Utility functions for UI helpers."""

from typing import Any

import pandas as pd
import pycountry
import streamlit as st


def init_session_state(key: str, default_value: Any) -> None:
    """
    Initialize session state key if not exists.

    Args:
        key: Session state key name
        default_value: Default value to set
    """
    if key not in st.session_state:
        st.session_state[key] = default_value


def format_datetime(dt: pd.Timestamp | None) -> str:
    """
    Format datetime for display.

    Args:
        dt: Datetime to format

    Returns:
        Formatted datetime string or empty string
    """
    if pd.isna(dt) or dt is None:
        return ""
    return dt.strftime("%Y-%m-%d %H:%M:%S")


def create_pagination_controls(
    total_rows: int, page_size: int = 100, key_prefix: str = ""
) -> tuple[int, int]:
    """
    Create pagination controls and return limit/offset.

    Args:
        total_rows: Total number of rows
        page_size: Number of rows per page
        key_prefix: Prefix for session state keys to avoid conflicts

    Returns:
        Tuple of (limit, offset) for query
    """
    # Initialize current page in session state
    page_key = f"{key_prefix}current_page"
    if page_key not in st.session_state:
        st.session_state[page_key] = 0

    total_pages = max(1, (total_rows + page_size - 1) // page_size)
    current_page = st.session_state[page_key]

    # Ensure current page is valid
    if current_page >= total_pages:
        current_page = total_pages - 1
        st.session_state[page_key] = current_page

    # Create pagination UI
    col1, col2, col3 = st.columns([1, 3, 1])

    with col1:
        if st.button("← Previous", disabled=current_page == 0):
            st.session_state[page_key] = max(0, current_page - 1)
            st.rerun()

    with col2:
        st.write(
            f"Page {current_page + 1} of {total_pages} "
            f"(Showing {page_size} rows per page)"
        )

    with col3:
        if st.button("Next →", disabled=current_page >= total_pages - 1):
            st.session_state[page_key] = min(total_pages - 1, current_page + 1)
            st.rerun()

    offset = current_page * page_size
    return page_size, offset


def validate_country_code(code: str) -> tuple[bool, str | None]:
    """
    Validate ISO Alpha-2 country code using pycountry.

    Args:
        code: Country code to validate (e.g., 'US', 'GB', 'JP')

    Returns:
        Tuple of (is_valid: bool, country_name: str | None)
        - is_valid: True if code is valid ISO Alpha-2 country code
        - country_name: Official country name if valid, None otherwise

    Examples:
        >>> validate_country_code("US")
        (True, "United States")
        >>> validate_country_code("GB")
        (True, "United Kingdom")
        >>> validate_country_code("XX")
        (False, None)
    """
    if not code or len(code) != 2 or not code.isalpha() or not code.isupper():
        return False, None

    try:
        country = pycountry.countries.get(alpha_2=code)
        if country:
            return True, country.name
        return False, None
    except (KeyError, AttributeError):
        return False, None


def search_countries(query: str) -> list[tuple[str, str]]:
    """
    Search for countries by name (partial match).

    Args:
        query: Search query (country name or partial name)

    Returns:
        List of tuples (country_code, country_name) matching the query
        Sorted by country name

    Examples:
        >>> search_countries("united")
        [('AE', 'United Arab Emirates'), ('GB', 'United Kingdom'), ('US', 'United States')]
        >>> search_countries("ukr")
        [('UA', 'Ukraine')]
    """
    if not query or len(query) < 2:
        return []

    query_lower = query.lower()
    results = []

    for country in pycountry.countries:
        # Search in official name
        if query_lower in country.name.lower():
            results.append((country.alpha_2, country.name))

    # Sort by name
    results.sort(key=lambda x: x[1])

    return results[:20]  # Limit to top 20 results