import logging
from typing import Any, Callable, List, Optional, Union
import pandas as pd
from utils.utils import colorText, get_sanitized_input

logger = logging.getLogger(__name__)

class Selector:
    @staticmethod
    def _get_sorted_items(items: List[Any], label_func: Callable[[Any], str]) -> List[Any]:
        return sorted(items, key=lambda item: label_func(item).lower())

    @staticmethod
    def _display_choices(
        items: List[Any],
        label_func: Callable[[Any], str],
        num_columns: int = 4,
        header: str = "Available Choices:"
    ) -> None:
        # Force single column if items are DataFrame rows
        
        if items and isinstance(items[0], (pd.Series, dict)):
            num_columns = 1

        rows = (len(items) + num_columns - 1) // num_columns
        print(f"\n{header}")
        for row in range(rows):
            line = ""
            for col in range(num_columns):
                idx = row + col * rows
                if idx < len(items):
                    label = label_func(items[idx])
                    line += f"{idx + 1}: {label:<30}"
            print(line)

    @staticmethod
    def _display_selected_items(
        selected: List[Any],
        label_func: Callable[[Any], str],
        num_columns: int = 4
    ) -> None:
        print(colorText("\nCurrent selections:", "cyan"))
        if not selected:
            print(" (none)")
            return
        sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
        rows = (len(sorted_selected) + num_columns - 1) // num_columns
        for row in range(rows):
            line = ""
            for col in range(num_columns):
                idx = row + col * rows
                if idx < len(sorted_selected):
                    label = label_func(sorted_selected[idx])
                    line += f"{label:<30}"
            print(line)

    @staticmethod
    def _parse_selection_input(input_str: str, max_index: int) -> List[int]:
        selections = []
        for part in input_str.split(","):
            part = part.strip()
            if "-" in part:
                try:
                    start, end = map(int, part.split("-"))
                    selections.extend(range(start, end + 1))
                except ValueError:
                    continue
            elif part.isdigit():
                selections.append(int(part))
        return [i for i in selections if 1 <= i <= max_index]

    @staticmethod
    def _select_from_list(
        items: List[Any],
        label_func: Callable[[Any], str],
        allow_multiple: bool = False,
        prompt_each: bool = False,
        header: str = "Available Choices:",
        num_columns: int = 4
    ) -> Union[Optional[Any], List[Any]]:
        if not items:
            logger.warning("No items available for selection.")
            return None

        full_sorted_items = Selector._get_sorted_items(items, label_func)
        remaining_items = full_sorted_items.copy()
        selected = []

        if allow_multiple:
            while True:
                Selector._display_choices(remaining_items, label_func, num_columns=num_columns, header=header)
                Selector._display_selected_items(selected, label_func, num_columns=num_columns)
                choice = get_sanitized_input("Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: ").strip().lower()
                if choice == "q":
                    break
                elif choice == "r":
                    selected.clear()
                    remaining_items = full_sorted_items.copy()
                    print(colorText("🔄 Selections reset.", "yellow"))
                    continue
                indices = Selector._parse_selection_input(choice, len(remaining_items))
                newly_selected = []
                for index in indices:
                    item = remaining_items[index - 1]
                    if item not in selected:
                        selected.append(item)
                        newly_selected.append(item)
                        if prompt_each:
                            logger.info(f"Selected: {label_func(item)}")
                    else:
                        logger.warning("Item already selected.")
                remaining_items = [item for item in remaining_items if item not in newly_selected]
            return selected if selected else None
        else:
            Selector._display_choices(full_sorted_items, label_func,  num_columns=num_columns, header=header)
            try:
                choice = int(get_sanitized_input("Select one item by number: "))
                if 1 <= choice <= len(full_sorted_items):
                    selected_item = full_sorted_items[choice - 1]
                    logger.info(f"Selected: {label_func(selected_item)}")
                    return selected_item
                else:
                    logger.warning("Selection out of range.")
            except ValueError:
                logger.warning("Invalid input.")
            return None

    @staticmethod
    def select_with_mode(
        items: List[Any],
        label_func: Callable[[Any], str],
        header: str = "Available Choices:"
    ) -> List[Any]:
        print(colorText("Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white"))
        mode = get_sanitized_input("").strip().lower()
        if mode == "a":
            return items
        selected = Selector._select_from_list(
            items,
            label_func=label_func,
            allow_multiple=True,
            prompt_each=False,
            header=header
        )
        if not selected:
            return items
        if mode == "i":
            print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
            return selected
        elif mode == "e":
            print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
            return [item for item in items if item not in selected]
        else:
            print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
            return items

    @staticmethod
    def select_objects(
        objects: List[Any],
        allow_multiple: bool = False,
        prompt_each: bool = False
    ) -> Union[Optional[Any], List[Any]]:
        return Selector._select_from_list(
            objects,
            label_func=lambda obj: getattr(obj, "name", str(obj)),
            allow_multiple=allow_multiple,
            prompt_each=prompt_each,
            header="Available Objects:"
        )

    @staticmethod
    def select_string(
        options: List[str],
        allow_multiple: bool = False,
        prompt_each: bool = False
    ) -> Union[Optional[str], List[str]]:
        return Selector._select_from_list(
            options,
            label_func=str,
            allow_multiple=allow_multiple,
            prompt_each=prompt_each,
            header="Available Options:"
        )

    @staticmethod
    def select_int(
        options: List[int],
        allow_multiple: bool = False,
        prompt_each: bool = False
    ) -> Union[Optional[int], List[int]]:
        return Selector._select_from_list(
            options,
            label_func=lambda x: str(x),
            allow_multiple=allow_multiple,
            prompt_each=prompt_each,
            header="Available Integers:"
        )

    @staticmethod
    def select_value(
        prompt: str,
        value_type: type = int,
        valid_range: Optional[tuple] = None,
        allow_quit: bool = False
    ) -> Optional[Any]:
        while True:
            user_input = get_sanitized_input(prompt).strip().lower()
            if allow_quit and user_input == "q":
                logger.info("User opted to quit value selection.")
                return None
            try:
                value = value_type(user_input)
                if valid_range:
                    min_val, max_val = valid_range
                    if not (min_val <= value <= max_val):
                        logger.warning(f"Value out of range ({min_val}–{max_val}).")
                        continue
                logger.info(f"User selected value: {value}")
                return value
            except ValueError:
                logger.warning(f"Invalid input. Expected a {value_type.__name__}.")

    @staticmethod
    def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
        while True:
            response = get_sanitized_input(prompt).strip().lower()
            if response in ["y", "yes"]:
                logger.info("User confirmed action.")
                return True
            elif response in ["n", "no"]:
                logger.info("User declined action.")
                return False
            else:
                logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")

    @staticmethod
    def select_dataframe_rows(
        df: pd.DataFrame,
        columns: Optional[List[str]] = None,
        allow_multiple: bool = False,
        prompt_each: bool = False,
        header: str = "Available Rows:"
    ) -> List[pd.Series]:
        if df.empty:
            print("DataFrame is empty.")
            return []

        if columns:
            df = df[columns]

        items = [row for _, row in df.iterrows()]
        label_func = lambda row: str(row.to_dict())

        result = Selector._select_from_list(
            items,
            label_func=label_func,
            allow_multiple=allow_multiple,
            prompt_each=prompt_each,
            header=header
        )

        if isinstance(result, pd.Series):
            return [result]
        elif isinstance(result, list):
            return result
        else:
            return []

    @staticmethod
    def select_dataframe_with_mode(
        df: pd.DataFrame,
        columns: Optional[List[str]] = None,
        header: str = "Available Rows:"
    ) -> List[pd.Series]:
        if df.empty:
            print("⚠️ DataFrame is empty.")
            return []

        # Filter columns if specified
        if columns:
            df = df[columns]

        items = df.to_dict("records")
        label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)

        # Show rows first
        print(colorText(header, "cyan"))
        for i, row in enumerate(items):
            print(f"{i}: {label_func(row)}")

        # Prompt for mode once
        print(colorText("\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):", "white"))
        mode = get_sanitized_input("").strip().lower()

        if mode == "a":
            return [pd.Series(row) for row in items]

        # Prompt for selection only once
        selected = Selector._select_from_list(
            items,
            label_func=label_func,
            allow_multiple=True,
            prompt_each=False,
            header=header
        )

        if not selected:
            return [pd.Series(row) for row in items]

        if mode == "i":
            print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
            return [pd.Series(row) for row in selected]
        elif mode == "e":
            print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
            return [pd.Series(row) for row in items if row not in selected]
        else:
            print(colorText("⚠️ Invalid mode. Returning all rows.", "yellow"))
            return [pd.Series(row) for row in items]