"""
Cohort Exploration CLI command for interactive cohort queries using the Cohort API.

BETA FEATURE

This command is in beta and subject to change or removal without notice:
- APIs may change in breaking ways between releases
- Functionality may be incomplete or unstable
- Access is limited to select users during the beta period

This command provides a streamlined interface for rapid validation of clinical hypotheses
via interactive prompts. Supports complex AND/OR/nested condition logic.

All models, API client, constants, and CLI helpers are self-contained in this module.
"""

from __future__ import annotations

import concurrent.futures
import json
import logging
import re
from datetime import date, datetime
from enum import Enum
from typing import (
    Annotated,
    Any,
    ClassVar,
    Dict,
    List,
    Literal,
    Optional,
    Tuple,
    Union,
)

import questionary
import typer
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator, validate_call
from pydantic import StrictFloat, StrictStr
from questionary import Choice, Style
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.tree import Tree
from typing_extensions import Self

from cli.util import ENVIRONMENT_OPTION, resolve_environment

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------


class SelectorType(str, Enum):
    """Available selector types for cohort queries."""

    MX_DIAGNOSIS_CODE = "mx_diagnosis_code"
    CPT_HCPCS_CODE = "cpt_hcpcs_code"
    NDC = "ndc"
    DRG_CODE = "drg_code"
    REVENUE_CODE = "revenue_code"

    @property
    def display_name(self) -> str:
        names = {
            SelectorType.MX_DIAGNOSIS_CODE: "Diagnosis Code (ICD-10)",
            SelectorType.CPT_HCPCS_CODE: "Procedure Code (CPT/HCPCS)",
            SelectorType.NDC: "Drug Reference (NDC)",
            SelectorType.DRG_CODE: "DRG Code",
            SelectorType.REVENUE_CODE: "Revenue Code",
        }
        return names.get(self, self.value)

    @property
    def examples(self) -> str:
        examples = {
            SelectorType.MX_DIAGNOSIS_CODE: "C34%, C50.1, Z51.11",
            SelectorType.CPT_HCPCS_CODE: "96413, 99213, J9271",
            SelectorType.NDC: "12345678901, 12345-6789-01",
            SelectorType.DRG_CODE: "470, 871",
            SelectorType.REVENUE_CODE: "0636, 0250",
        }
        return examples.get(self, "")

    @property
    def event_table(self) -> str:
        tables = {
            SelectorType.MX_DIAGNOSIS_CODE: "mx",
            SelectorType.CPT_HCPCS_CODE: "mx",
            SelectorType.NDC: "rx",
            SelectorType.DRG_CODE: "mx",
            SelectorType.REVENUE_CODE: "mx",
        }
        return tables.get(self, "mx")

    @property
    def event_column(self) -> str:
        columns = {
            SelectorType.MX_DIAGNOSIS_CODE: "header_diagnosis_code",
            SelectorType.CPT_HCPCS_CODE: "cpt_hcpcs_code",
            SelectorType.NDC: "ndc",
            SelectorType.DRG_CODE: "drg_code",
            SelectorType.REVENUE_CODE: "revenue_code",
        }
        return columns.get(self, self.value)


class CountType(str, Enum):
    """Type of patient count to retrieve."""

    APPROXIMATE = "approximate"
    DISTINCT = "distinct"

    @property
    def display_name(self) -> str:
        names = {
            CountType.APPROXIMATE: "Approximate (fast)",
            CountType.DISTINCT: "Distinct (exact)",
        }
        return names.get(self, self.value)


# ---------------------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------------------

CODE_VALIDATION_PATTERNS = {
    SelectorType.MX_DIAGNOSIS_CODE: r"^[A-Za-z]\d{1,6}[A-Za-z]?%?$",
    SelectorType.CPT_HCPCS_CODE: r"^(\d{3,5}|[A-Za-z]\d{1,4})%?$",
    SelectorType.NDC: r"^\d{5,13}%?$",
    SelectorType.DRG_CODE: r"^\d{1,4}%?$",
    SelectorType.REVENUE_CODE: r"^\d{1,4}%?$",
}


class DateRange(BaseModel):
    """A date range for filtering cohort queries."""

    from_date: date
    to_date: date

    model_config = ConfigDict(
        populate_by_name=True,
        validate_assignment=True,
        protected_namespaces=(),
        extra="forbid",
    )

    @field_validator("from_date", "to_date", mode="before")
    @classmethod
    def parse_date(cls, v: Union[str, date, datetime]) -> date:
        if isinstance(v, datetime):
            return v.date()
        if isinstance(v, date):
            return v
        if isinstance(v, str):
            try:
                return datetime.strptime(v, "%Y-%m-%d").date()
            except ValueError:
                pass
            try:
                return datetime.strptime(v, "%m/%d/%y").date()
            except ValueError:
                pass
            raise ValueError(f"Invalid date format: {v}. Use YYYY-MM-DD or mm/dd/yy")
        raise ValueError(f"Cannot parse date from {type(v)}")

    @classmethod
    def from_strings(
        cls,
        from_date: str,
        to_date: Optional[str] = None,
        format: str = "%m/%d/%y",
    ) -> "DateRange":
        parsed_from = datetime.strptime(from_date, format).date()
        parsed_to = datetime.strptime(to_date, format).date() if to_date else date.today()
        return cls(from_date=parsed_from, to_date=parsed_to)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "from": self.from_date.isoformat(),
            "to": self.to_date.isoformat(),
        }


class Condition(BaseModel):
    """A single condition with a selector type and one or more values.

    Multiple values within a single condition are combined with OR logic.
    """

    selector_type: SelectorType
    values: Annotated[List[str], Field(min_length=1)]

    model_config = ConfigDict(
        populate_by_name=True,
        validate_assignment=True,
        protected_namespaces=(),
        extra="forbid",
    )

    @model_validator(mode="before")
    @classmethod
    def normalize_and_validate_values(cls, data: Any) -> Any:
        if not isinstance(data, dict):
            return data

        selector_type = data.get("selector_type")
        values = data.get("values")
        if selector_type is None or values is None:
            return data

        if isinstance(selector_type, str):
            try:
                selector_type_enum = SelectorType(selector_type)
            except ValueError:
                return data
        else:
            selector_type_enum = selector_type

        if isinstance(values, str):
            values = [val.strip() for val in values.split(",") if val.strip()]

        normalized = []
        pattern = CODE_VALIDATION_PATTERNS.get(selector_type_enum)
        for val in values:
            normalized_val = val
            if selector_type_enum == SelectorType.NDC:
                normalized_val = val.replace("-", "")
            elif selector_type_enum == SelectorType.MX_DIAGNOSIS_CODE:
                normalized_val = val.replace(".", "")
            if pattern and not re.match(pattern, normalized_val):
                raise ValueError(f"Invalid {selector_type_enum.display_name} format: '{val}'. " f"Examples: {selector_type_enum.examples}")
            normalized.append(normalized_val)

        data["values"] = normalized
        return data

    def to_logic(self) -> Dict[str, Any]:
        if len(self.values) == 1:
            return {self.selector_type.value: self.values[0]}
        return {self.selector_type.value: self.values}

    def to_events(self) -> Dict[str, Any]:
        table = self.selector_type.event_table
        col = self.selector_type.event_column
        if len(self.values) == 1:
            return {
                "table": table,
                "filter": {"col": col, "op": "eq", "value": self.values[0]},
            }
        return {
            "op": "or",
            "conditions": [{"table": table, "filter": {"col": col, "op": "eq", "value": v}} for v in self.values],
        }

    def describe(self) -> str:
        if len(self.values) == 1:
            return f"{self.selector_type.display_name}: {self.values[0]}"
        return f"{self.selector_type.display_name}: {' OR '.join(self.values)}"


class ConditionGroup(BaseModel):
    """A group of conditions combined with AND or OR logic."""

    operator: Literal["and", "or"]
    items: List[Union[Condition, "ConditionGroup"]] = Field(default_factory=list)

    model_config = ConfigDict(
        populate_by_name=True,
        validate_assignment=True,
        protected_namespaces=(),
        extra="forbid",
    )

    def to_logic(self) -> Dict[str, Any]:
        if len(self.items) == 0:
            raise ValueError("Condition group cannot be empty")
        if len(self.items) == 1:
            return self.items[0].to_logic()
        return {
            "op": self.operator,
            "conditions": [item.to_logic() for item in self.items],
        }

    def to_events(self) -> Dict[str, Any]:
        if len(self.items) == 0:
            raise ValueError("Condition group cannot be empty")
        if len(self.items) == 1:
            return self.items[0].to_events()
        return {
            "op": self.operator,
            "conditions": [item.to_events() for item in self.items],
        }

    def describe(self, indent: int = 0) -> str:
        prefix = "  " * indent
        if len(self.items) == 0:
            return f"{prefix}(empty group)"
        if len(self.items) == 1:
            if isinstance(self.items[0], Condition):
                return f"{prefix}{self.items[0].describe()}"
            return self.items[0].describe(indent)
        lines = []
        for item in self.items:
            if isinstance(item, Condition):
                lines.append(f"{prefix}\u2022 {item.describe()}")
            else:
                lines.append(f"{prefix}\u2022 ({item.operator.upper()} group):")
                lines.append(item.describe(indent + 2))
        return "\n".join(lines)


ConditionGroup.model_rebuild()


# ---------------------------------------------------------------------------
# Logic-tree utilities
# ---------------------------------------------------------------------------


def logic_to_events(logic: Dict[str, Any]) -> Dict[str, Any]:
    """Convert a cohort logic tree to the events format used by distinct_patient_count."""
    if "op" in logic:
        return {
            "op": logic["op"],
            "conditions": [logic_to_events(cond) for cond in logic.get("conditions", [])],
        }
    for selector_key, value in logic.items():
        selector_type = SelectorType(selector_key)
        table = selector_type.event_table
        col = selector_type.event_column
        values = value if isinstance(value, list) else [value]
        if len(values) > 1:
            return {
                "op": "or",
                "conditions": [{"table": table, "filter": {"col": col, "op": "eq", "value": v}} for v in values],
            }
        return {"table": table, "filter": {"col": col, "op": "eq", "value": values[0]}}
    return {}


def describe_logic_tree(logic: Dict[str, Any], indent: int = 0) -> str:
    """Generate a human-readable description of a logic tree dictionary."""
    prefix = "  " * indent
    if "op" in logic:
        op = logic["op"].upper()
        conditions = logic.get("conditions", [])
        lines = [f"{prefix}({op})"]
        for cond in conditions:
            lines.append(describe_logic_tree(cond, indent + 1))
        return "\n".join(lines)
    else:
        for selector_key, value in logic.items():
            try:
                selector_type = SelectorType(selector_key)
                type_name = selector_type.display_name
            except ValueError:
                type_name = selector_key
            if isinstance(value, list):
                return f"{prefix}\u2022 {type_name}: {' OR '.join(value)}"
            else:
                return f"{prefix}\u2022 {type_name}: {value}"
    return f"{prefix}(unknown)"


# ---------------------------------------------------------------------------
# Response models
# ---------------------------------------------------------------------------


class SelectorPatientCountResponse(BaseModel):
    """Response from the selector-patient-count endpoint."""

    total_count: int
    warnings: List[str] = Field(default_factory=list)

    model_config = ConfigDict(
        populate_by_name=True,
        validate_assignment=True,
        protected_namespaces=(),
        extra="allow",
    )

    @classmethod
    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
        if obj is None:
            return None
        if not isinstance(obj, dict):
            return cls.model_validate(obj)
        return cls(
            total_count=obj.get("totalCount", 0),
            warnings=obj.get("warnings", []),
        )


class DistinctCountResponse(BaseModel):
    """Response from the distinct-patient-count endpoint."""

    total: int
    counts: Dict[str, int] = Field(default_factory=dict)
    warnings: List[str] = Field(default_factory=list)

    model_config = ConfigDict(
        populate_by_name=True,
        validate_assignment=True,
        protected_namespaces=(),
        extra="allow",
    )

    @classmethod
    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
        if obj is None:
            return None
        if not isinstance(obj, dict):
            return cls.model_validate(obj)
        return cls(
            total=obj.get("total", 0),
            counts=obj.get("counts", {}),
            warnings=obj.get("warnings", []),
        )


# ---------------------------------------------------------------------------
# Minimal Cohort API client
# ---------------------------------------------------------------------------

CONTENT_TYPE_JSON = "application/json"


class CohortApi:
    """Minimal Cohort API client for the cohort-exploration CLI.

    Only exposes the three methods needed by the interactive CLI:
    ``selector_patient_count``, ``distinct_patient_count``, and
    ``get_total_count``.
    """

    def __init__(self, api_client: Optional[Any] = None) -> None:
        from komodo.api_client import ApiClient

        if api_client is None:
            api_client = ApiClient.get_default()
        self.api_client = api_client
        self.api_version = "2024-10-01"

    def _get_host(self) -> str:
        from komodo.config import SDKSettings

        return SDKSettings().proxy_domain

    def _make_headers(self, extra: Optional[Dict[str, Any]] = None) -> Dict[str, Optional[str]]:
        headers: Dict[str, Optional[str]] = extra.copy() if extra else {}
        headers.setdefault("Accept", CONTENT_TYPE_JSON)
        headers.setdefault("Content-Type", CONTENT_TYPE_JSON)
        headers["X-API-Version"] = self.api_version
        return headers

    def selector_patient_count(
        self,
        logic: Dict[str, Any],
        patient_type: str = "adv",
        date_range: Optional[DateRange] = None,
    ) -> SelectorPatientCountResponse:
        """Get approximate patient count for a cohort query."""
        from komodo.exceptions import ApiException

        cohort_payload: Dict[str, Any] = {
            "logic": logic,
            "patientType": patient_type,
            "validateValues": False,
        }
        if date_range:
            cohort_payload["dateRange"] = date_range.to_dict()

        param = self.api_client.param_serialize(
            method="POST",
            resource_path="/burrow/v1/cohort/selector-patient-count",
            body={"cohort": cohort_payload},
            header_params=self._make_headers(),
            auth_settings=["HTTPBearer", "x-account-id"],
            _host=self._get_host(),
        )

        response_data = self.api_client.call_api(*param, _request_timeout=300.0)
        response_data.read()
        response_text = response_data.data.decode("utf-8")

        if not 200 <= response_data.status <= 299:
            raise ApiException.from_response(http_resp=response_data, body=response_text, data=None)

        return SelectorPatientCountResponse.from_dict(json.loads(response_text))

    def distinct_patient_count(
        self,
        logic: Dict[str, Any],
        date_range: DateRange,
        patient_type: str = "adv",
        granularity: Optional[Literal["year", "quarter", "month"]] = None,
        patient_criteria: Optional[Dict[str, Any]] = None,
    ) -> DistinctCountResponse:
        """Get distinct patient count for a cohort query."""
        from komodo.exceptions import ApiException

        events = logic_to_events(logic)
        payload: Dict[str, Any] = {
            "dateRange": date_range.to_dict(),
            "patientType": patient_type,
            "events": events,
        }
        if granularity:
            payload["granularity"] = granularity
        if patient_criteria:
            payload["patientCriteria"] = patient_criteria

        param = self.api_client.param_serialize(
            method="POST",
            resource_path="/burrow/v1/cohort/distinct-patient-count",
            body=payload,
            header_params=self._make_headers(),
            auth_settings=["HTTPBearer", "x-account-id"],
            _host=self._get_host(),
        )

        response_data = self.api_client.call_api(*param, _request_timeout=300.0)
        response_data.read()
        response_text = response_data.data.decode("utf-8")

        if not 200 <= response_data.status <= 299:
            raise ApiException.from_response(http_resp=response_data, body=response_text, data=None)

        return DistinctCountResponse.from_dict(json.loads(response_text))

    def get_total_count(
        self,
        logic: Dict[str, Any],
        count_type: CountType,
        date_range: Optional[DateRange] = None,
        patient_type: str = "adv",
    ) -> Tuple[int, Union[SelectorPatientCountResponse, DistinctCountResponse]]:
        """Get patient count using the specified count type."""
        if count_type == CountType.DISTINCT:
            if not date_range:
                raise ValueError("Date range is required for distinct patient counts")
            response = self.distinct_patient_count(logic=logic, date_range=date_range, patient_type=patient_type)
            return response.total, response
        else:
            response = self.selector_patient_count(logic=logic, date_range=date_range, patient_type=patient_type)
            return response.total_count, response


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

VALID_STATE_CODES = {
    "AL",
    "AK",
    "AZ",
    "AR",
    "CA",
    "CO",
    "CT",
    "DE",
    "FL",
    "GA",
    "HI",
    "ID",
    "IL",
    "IN",
    "IA",
    "KS",
    "KY",
    "LA",
    "ME",
    "MD",
    "MA",
    "MI",
    "MN",
    "MS",
    "MO",
    "MT",
    "NE",
    "NV",
    "NH",
    "NJ",
    "NM",
    "NY",
    "NC",
    "ND",
    "OH",
    "OK",
    "OR",
    "PA",
    "RI",
    "SC",
    "SD",
    "TN",
    "TX",
    "UT",
    "VT",
    "VA",
    "WA",
    "WV",
    "WI",
    "WY",
    "DC",
    "PR",
    "VI",
    "GU",
    "AS",
    "MP",
}

TOP_10_STATES = ["FL", "CA", "NY", "TX", "PA", "OH", "IL", "MI", "NC", "GA"]
AGE_BUCKETS = [
    ("0-17", 0, 17),
    ("18-34", 18, 34),
    ("35-44", 35, 44),
    ("45-54", 45, 54),
    ("55-64", 55, 64),
    ("65-74", 65, 74),
    ("75+", 75, 120),
]
RACE_ETHNICITY = [
    "WHITE",
    "BLACK OR AFRICAN AMERICAN",
    "HISPANIC OR LATINO",
    "ASIAN OR PACIFIC ISLANDER",
]

QUESTIONARY_STYLE = Style(
    [
        ("qmark", "fg:#00b7c2 bold"),
        ("question", "bold"),
        ("answer", "fg:#f44336 bold"),
        ("pointer", "fg:#00b7c2 bold"),
        ("highlighted", "fg:#00b7c2 bold"),
        ("selected", "fg:#00b7c2 bold"),
        ("separator", "fg:#cc5454"),
        ("instruction", "fg:#858585"),
        ("text", ""),
        ("disabled", "fg:#858585 italic"),
    ]
)


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------


def validate_code(code: str, selector_type: SelectorType) -> str | bool:
    """
    Validate a code against its selector type's pattern.

    Returns True if valid, or an error message string if invalid.
    Normalizes codes before validation (removes dots from diagnosis, hyphens from NDC).
    """
    if not code or not code.strip():
        return "At least one value is required"

    pattern = CODE_VALIDATION_PATTERNS.get(selector_type)
    if not pattern:
        return True

    test_code = code
    if selector_type == SelectorType.NDC:
        test_code = code.replace("-", "")
    elif selector_type == SelectorType.MX_DIAGNOSIS_CODE:
        test_code = code.replace(".", "")

    if not re.match(pattern, test_code):
        return f"Invalid format for '{code}'. Examples: {selector_type.examples}"
    return True


def validate_codes(codes_str: str, selector_type: SelectorType) -> str | bool:
    """Validate comma-separated codes for a selector type."""
    if not codes_str or not codes_str.strip():
        return True

    values = [v.strip() for v in codes_str.split(",") if v.strip()]
    for val in values:
        result = validate_code(val, selector_type)
        if result is not True:
            return result
    return True


def validate_date_format(date_str: str) -> bool:
    """Validate date in mm/dd/yy format."""
    if not date_str:
        return True
    try:
        datetime.strptime(date_str, "%m/%d/%y")
        return True
    except ValueError:
        return False


def parse_values_input(values_input: str) -> list[str]:
    """Parse comma-separated user input into a list of values."""
    return [val.strip() for val in values_input.split(",") if val.strip()]


def format_values_for_display(values: list[str]) -> str:
    """Format values for display, using OR for multiple values."""
    if len(values) == 1:
        return values[0]
    return " OR ".join(values)


def describe_condition_for_display(
    condition: Condition,
    display_overrides: Optional[dict[int, list[str]]] = None,
) -> str:
    """Describe a condition, using dotted ICD-10 values when available."""
    if display_overrides and condition.selector_type == SelectorType.MX_DIAGNOSIS_CODE:
        display_values = display_overrides.get(id(condition))
        if display_values:
            return f"{condition.selector_type.display_name}: {format_values_for_display(display_values)}"
    return condition.describe()


def describe_group_for_display(
    group: ConditionGroup,
    display_overrides: Optional[dict[int, list[str]]] = None,
    indent: int = 0,
) -> str:
    """Describe a condition group with optional ICD-10 display overrides."""
    prefix = "  " * indent
    if len(group.items) == 0:
        return f"{prefix}(empty group)"
    if len(group.items) == 1:
        item = group.items[0]
        if isinstance(item, Condition):
            return f"{prefix}{describe_condition_for_display(item, display_overrides)}"
        return describe_group_for_display(item, display_overrides, indent)

    lines = []
    for item in group.items:
        if isinstance(item, Condition):
            lines.append(f"{prefix}\u2022 {describe_condition_for_display(item, display_overrides)}")
        else:
            lines.append(f"{prefix}\u2022 ({item.operator.upper()} group):")
            lines.append(describe_group_for_display(item, display_overrides, indent + 2))
    return "\n".join(lines)


def build_condition_from_input(
    selector_type: SelectorType,
    values_input: str,
    display_overrides: Optional[dict[int, list[str]]] = None,
) -> Condition:
    """Create a condition and capture dotted ICD-10 values for display."""
    condition = Condition(selector_type=selector_type, values=values_input)
    if display_overrides is not None and selector_type == SelectorType.MX_DIAGNOSIS_CODE:
        display_overrides[id(condition)] = parse_values_input(values_input)
    return condition


def display_current_conditions(
    console: Console,
    conditions: list,
    operator: str,
    display_overrides: Optional[dict[int, list[str]]] = None,
) -> None:
    """Display the current conditions in a formatted tree view."""
    if not conditions:
        console.print(Text("  No conditions added yet.", style="dim"))
        return

    tree = Tree(f"[bold cyan]Query ({operator.upper()} logic)[/bold cyan]")
    for i, item in enumerate(conditions, 1):
        if isinstance(item, Condition):
            tree.add(f"[green]{i}.[/green] {describe_condition_for_display(item, display_overrides)}")
        else:
            group_branch = tree.add(f"[green]{i}.[/green] [yellow]({item.operator.upper()} group)[/yellow]")
            for j, sub_item in enumerate(item.items, 1):
                if isinstance(sub_item, Condition):
                    group_branch.add(f"[dim]{j}.[/dim] {describe_condition_for_display(sub_item, display_overrides)}")

    console.print(tree)


# ---------------------------------------------------------------------------
# Interactive prompts
# ---------------------------------------------------------------------------


def prompt_add_condition(
    console: Console,
    display_overrides: Optional[dict[int, list[str]]] = None,
) -> Optional[Condition]:
    """Prompt user to add a single condition."""
    choices = [Choice(f"{st.display_name} (e.g., {st.examples})", st) for st in SelectorType]

    selector_type: SelectorType | None = questionary.select(
        "Select condition type:",
        choices=choices,
        style=QUESTIONARY_STYLE,
    ).ask()

    if selector_type is None:
        return None

    console.print(Text("  \U0001f4a1 Tip: Enter multiple values separated by commas for OR logic (e.g., C34%,C50%)", style="dim"))
    console.print(Text("  \U0001f4a1 Use % suffix for prefix matching (e.g., C34% matches C34, C34.1, C34.9)", style="dim"))

    values_input = questionary.text(
        f"{selector_type.display_name}:",
        validate=lambda x: validate_codes(x, selector_type) if x.strip() else "At least one value is required",
        style=QUESTIONARY_STYLE,
    ).ask()

    if values_input is None:
        return None

    try:
        return build_condition_from_input(selector_type, values_input, display_overrides)
    except ValueError as e:
        console.print(Text(f"  \u274c {e}", style="bold red"))
        return None


def prompt_create_group(
    console: Console,
    display_overrides: Optional[dict[int, list[str]]] = None,
) -> Optional[ConditionGroup]:
    """Prompt user to create a nested condition group."""
    console.print(
        Panel(
            "[bold]Create a Condition Group[/bold]\n\n" "A group combines multiple conditions with AND or OR logic.\n" "Groups can be nested within the main query.",
            border_style="yellow",
        )
    )

    group_operator = questionary.select(
        "How should conditions in this group be combined?",
        choices=[
            Choice("AND - patients must match ALL conditions", "and"),
            Choice("OR - patients must match ANY condition", "or"),
        ],
        style=QUESTIONARY_STYLE,
    ).ask()

    if group_operator is None:
        return None

    group = ConditionGroup(operator=group_operator)

    while True:
        console.print()
        console.print(Text(f"  Group conditions ({group_operator.upper()} logic):", style="bold yellow"))
        if group.items:
            for i, item in enumerate(group.items, 1):
                console.print(Text(f"    {i}. {describe_condition_for_display(item, display_overrides)}", style="cyan"))
        else:
            console.print(Text("    (none yet)", style="dim"))

        action = questionary.select(
            "Group action:",
            choices=[
                Choice("\u2795 Add condition to group", "add"),
                Choice("\u2705 Finish group" + (" (need at least 1)" if not group.items else ""), "done"),
                Choice("\u274c Cancel group", "cancel"),
            ],
            style=QUESTIONARY_STYLE,
        ).ask()

        if action is None or action == "cancel":
            return None
        elif action == "done":
            if not group.items:
                console.print(Text("  \u26a0\ufe0f  Group must have at least one condition.", style="yellow"))
                continue
            return group
        elif action == "add":
            condition = prompt_add_condition(console, display_overrides)
            if condition:
                group.items.append(condition)


def validate_state_codes(codes_str: str) -> str | bool:
    """Validate comma-separated state codes against valid US state/territory codes."""
    if not codes_str or not codes_str.strip():
        return True

    values = [v.strip().upper() for v in codes_str.split(",") if v.strip()]
    for val in values:
        if val not in VALID_STATE_CODES:
            return f"Invalid state code: '{val}'. Use valid 2-letter US state codes (e.g., AZ, NV, WA)"
    return True


def prompt_for_additional_states(console: Console) -> list[str]:
    """Prompt user to optionally add specific state codes beyond the top 10."""
    console.print()
    console.print(Text(f"  \U0001f4cd Default states queried: {', '.join(TOP_10_STATES)}", style="dim"))

    add_states = questionary.confirm(
        "Add specific state codes to query?",
        default=False,
        style=QUESTIONARY_STYLE,
    ).ask()

    if not add_states:
        return []

    console.print(Text("  \U0001f4a1 Enter 2-letter state codes separated by commas (e.g., AZ, NV, WA, CO)", style="dim"))

    states_input = questionary.text(
        "Additional State Codes:",
        validate=validate_state_codes,
        style=QUESTIONARY_STYLE,
    ).ask()

    if not states_input or not states_input.strip():
        return []

    additional_states = [s.strip().upper() for s in states_input.split(",") if s.strip()]
    additional_states = [s for s in additional_states if s not in TOP_10_STATES]
    seen = set()
    unique_states = []
    for s in additional_states:
        if s not in seen:
            seen.add(s)
            unique_states.append(s)

    if unique_states:
        console.print(Text(f"  \u2705 Added states: {', '.join(unique_states)}", style="green"))

    return unique_states


def prompt_for_date_range(console: Console, required: bool = True) -> tuple[Optional[str], Optional[str]]:
    """Prompt for date range (required for comprehensive summary)."""
    console.print()

    if not required:
        add_dates = questionary.confirm(
            "Add date range filter?",
            default=False,
            style=QUESTIONARY_STYLE,
        ).ask()

        if not add_dates:
            return None, None
    else:
        console.print(Text("  \U0001f4c5 Date range is required for the comprehensive summary.", style="dim"))

    def validate_required_date(x):
        if not x or not x.strip():
            return "Date is required"
        if not validate_date_format(x):
            return "Invalid date format. Use mm/dd/yy"
        return True

    date_from = questionary.text(
        "Start Date (mm/dd/yy, e.g., 01/01/23):",
        validate=validate_required_date if required else lambda x: validate_date_format(x) or "Invalid date format. Use mm/dd/yy",
        style=QUESTIONARY_STYLE,
    ).ask()

    if date_from is None:
        return None, None

    date_to = questionary.text(
        "End Date (mm/dd/yy, e.g., 12/31/25, leave empty for today):",
        validate=lambda x: (not x) or validate_date_format(x) or "Invalid date format. Use mm/dd/yy",
        style=QUESTIONARY_STYLE,
    ).ask()

    if not date_to or not date_to.strip():
        date_to = datetime.now().strftime("%m/%d/%y")

    return date_from.strip() if date_from else None, date_to.strip() if date_to else None


def prompt_for_input(console: Console) -> dict:
    """Interactively prompt the user for cohort parameters with support for complex logic."""
    console.print(
        Panel(
            "[bold cyan]Cohort Exploration Tool[/bold cyan]\n\n"
            "Build complex cohort queries with AND/OR logic.\n"
            "\u2022 Add conditions one at a time\n"
            "\u2022 Combine conditions with AND or OR\n"
            "\u2022 Create nested groups for complex logic\n"
            "\u2022 Use [bold]%[/bold] suffix for prefix matching (e.g., C34% matches C34, C34.1, C34.9)",
            title="\U0001f52c Cohort API Query Builder",
            border_style="cyan",
        )
    )

    mode = questionary.select(
        "Query mode:",
        choices=[
            Choice("Simple - guided prompts for common fields", "simple"),
            Choice("Advanced - build complex AND/OR logic", "advanced"),
        ],
        style=QUESTIONARY_STYLE,
    ).ask()

    if mode is None:
        raise typer.Exit(code=0)

    if mode == "simple":
        return prompt_simple_mode(console)
    else:
        return prompt_advanced_mode(console)


def prompt_simple_mode(console: Console) -> dict:
    """Simple guided prompt mode."""
    console.print(Text("\n  \U0001f4a1 Enter comma-separated values for OR logic within each field.", style="dim"))
    console.print(Text("  \U0001f4a1 All fields are combined with AND logic.", style="dim"))
    console.print()

    diagnosis_input = questionary.text(
        "Diagnosis Code(s) (ICD-10, e.g., C34%,C50%):",
        validate=lambda x: validate_codes(x, SelectorType.MX_DIAGNOSIS_CODE),
        style=QUESTIONARY_STYLE,
    ).ask()

    if diagnosis_input is None:
        raise typer.Exit(code=0)

    procedure_input = questionary.text(
        "Procedure Code(s) (CPT/HCPCS, e.g., 96413,96415):",
        validate=lambda x: validate_codes(x, SelectorType.CPT_HCPCS_CODE),
        style=QUESTIONARY_STYLE,
    ).ask()

    if procedure_input is None:
        raise typer.Exit(code=0)

    ndc_input = questionary.text(
        "Drug Reference(s) (NDC, e.g., 12345678901):",
        validate=lambda x: validate_codes(x, SelectorType.NDC),
        style=QUESTIONARY_STYLE,
    ).ask()

    if ndc_input is None:
        raise typer.Exit(code=0)

    conditions: list[Condition] = []
    display_overrides: dict[int, list[str]] = {}

    if diagnosis_input and diagnosis_input.strip():
        conditions.append(build_condition_from_input(SelectorType.MX_DIAGNOSIS_CODE, diagnosis_input, display_overrides))

    if procedure_input and procedure_input.strip():
        conditions.append(build_condition_from_input(SelectorType.CPT_HCPCS_CODE, procedure_input, display_overrides))

    if ndc_input and ndc_input.strip():
        conditions.append(build_condition_from_input(SelectorType.NDC, ndc_input, display_overrides))

    if len(conditions) < 1:
        console.print(Text("\u274c At least one condition is required.", style="bold red"))
        raise typer.Exit(code=1)

    root_group = ConditionGroup(operator="and", items=conditions)

    console.print()
    display_current_conditions(console, conditions, "and", display_overrides)

    output_mode = prompt_output_mode(console)
    count_type = prompt_count_type(console, output_mode)

    additional_states = []
    if output_mode == "comprehensive":
        additional_states = prompt_for_additional_states(console)

    date_required = output_mode == "comprehensive" or count_type == CountType.DISTINCT
    date_from, date_to = prompt_for_date_range(console, required=date_required)

    return {
        "logic": root_group.to_logic(),
        "date_from": date_from,
        "date_to": date_to,
        "output_mode": output_mode,
        "count_type": count_type,
        "additional_states": additional_states,
    }


def prompt_output_mode(console: Console) -> str:
    """Prompt user to choose between quick count or comprehensive summary."""
    console.print()
    output_mode = questionary.select(
        "What output would you like?",
        choices=[
            Choice("\U0001f4ca Quick count - just the total patient count (faster)", "quick"),
            Choice("\U0001f4c8 Comprehensive summary - demographics, trends, bar charts", "comprehensive"),
        ],
        style=QUESTIONARY_STYLE,
    ).ask()

    if output_mode is None:
        raise typer.Exit(code=0)

    return output_mode


def prompt_count_type(console: Console, output_mode: str) -> CountType:
    """Prompt user to choose between approximate or distinct patient counts."""
    console.print()

    if output_mode == "comprehensive":
        console.print(Text("  \U0001f4a1 Comprehensive summary uses distinct counts for the total and all breakdowns.", style="dim"))

    count_type = questionary.select(
        "Count type for total patients:",
        choices=[
            Choice(f"\u26a1 {CountType.APPROXIMATE.display_name}", CountType.APPROXIMATE),
            Choice(f"\U0001f3af {CountType.DISTINCT.display_name}", CountType.DISTINCT),
        ],
        style=QUESTIONARY_STYLE,
    ).ask()

    if count_type is None:
        raise typer.Exit(code=0)

    return count_type


def prompt_advanced_mode(console: Console) -> dict:
    """Advanced mode for building complex AND/OR logic."""
    conditions: list[Condition | ConditionGroup] = []
    root_operator = "and"
    display_overrides: dict[int, list[str]] = {}

    while True:
        console.print()
        display_current_conditions(console, conditions, root_operator, display_overrides)
        console.print()

        action_choices = [
            Choice("\u2795 Add condition", "add"),
            Choice("\U0001f4c1 Create nested group (AND/OR)", "group"),
        ]

        if len(conditions) >= 2:
            action_choices.append(Choice(f"\U0001f500 Change root operator (currently: {root_operator.upper()})", "operator"))

        if conditions:
            action_choices.append(Choice("\U0001f5d1\ufe0f  Remove last condition", "remove"))
            action_choices.append(Choice("\u2705 Build and run query", "done"))

        action_choices.append(Choice("\u274c Cancel", "cancel"))

        action = questionary.select(
            "What would you like to do?",
            choices=action_choices,
            style=QUESTIONARY_STYLE,
        ).ask()

        if action is None or action == "cancel":
            raise typer.Exit(code=0)

        elif action == "add":
            condition = prompt_add_condition(console, display_overrides)
            if condition:
                conditions.append(condition)

        elif action == "group":
            group = prompt_create_group(console, display_overrides)
            if group:
                conditions.append(group)

        elif action == "operator":
            new_op = questionary.select(
                "Select root operator:",
                choices=[
                    Choice("AND - patients must match ALL conditions", "and"),
                    Choice("OR - patients must match ANY condition", "or"),
                ],
                style=QUESTIONARY_STYLE,
            ).ask()
            if new_op:
                root_operator = new_op

        elif action == "remove":
            if conditions:
                removed = conditions.pop()
                if isinstance(removed, Condition):
                    console.print(Text(f"  Removed: {describe_condition_for_display(removed, display_overrides)}", style="yellow"))
                else:
                    console.print(Text(f"  Removed: ({removed.operator.upper()} group)", style="yellow"))

        elif action == "done":
            break

    if not conditions:
        console.print(Text("\u274c At least one condition is required.", style="bold red"))
        raise typer.Exit(code=1)

    root_group = ConditionGroup(operator=root_operator, items=conditions)

    console.print()
    console.print(
        Panel(
            describe_group_for_display(root_group, display_overrides),
            title="\U0001f4cb Final Query",
            border_style="green",
        )
    )

    output_mode = prompt_output_mode(console)
    count_type = prompt_count_type(console, output_mode)

    additional_states = []
    if output_mode == "comprehensive":
        additional_states = prompt_for_additional_states(console)

    date_required = output_mode == "comprehensive" or count_type == CountType.DISTINCT
    date_from, date_to = prompt_for_date_range(console, required=date_required)

    return {
        "logic": root_group.to_logic(),
        "date_from": date_from,
        "date_to": date_to,
        "output_mode": output_mode,
        "count_type": count_type,
        "additional_states": additional_states,
    }


# ---------------------------------------------------------------------------
# Display helpers
# ---------------------------------------------------------------------------


def make_bar(value: int, max_value: int, width: int = 40) -> str:
    """Create a text-based bar chart segment, padded to fixed width for alignment."""
    if max_value == 0:
        return " " * width
    bar_len = int((value / max_value) * width)
    bar = "\u2588" * bar_len
    return bar.ljust(width)


def format_number(n: int) -> str:
    """Format number with commas."""
    return f"{n:,}"


def format_percent(value: int, total: int) -> str:
    """Format percentage with one decimal place."""
    if total == 0:
        return "0.0%"
    return f"{(value / total) * 100:.1f}%"


def run_comprehensive_summary_streaming(
    cohort_api: CohortApi,
    logic: dict,
    date_range: DateRange,
    console: Console,
    additional_states: Optional[list[str]] = None,
    count_type: CountType = CountType.APPROXIMATE,
) -> None:
    """Run breakdown queries in parallel and print results as they arrive.

    Comprehensive mode always uses distinct counts for both the total and all
    breakdowns (sex, state, age, race, deceased, time) so that percentages and
    totals are internally consistent. Approximate is not supported for breakdowns
    (no API for approx + patient_criteria/granularity), so we force distinct here.
    """
    # Force distinct so total_adv and all breakdown queries use the same count type.
    count_type = CountType.DISTINCT

    results: dict = {}
    errors: dict[str, Exception] = {}
    query_count = 0
    start_time = datetime.now()

    all_states = list(TOP_10_STATES)
    if additional_states:
        all_states.extend(additional_states)

    def get_count_for_type(patient_type: str) -> int:
        if count_type == CountType.DISTINCT:
            response = cohort_api.distinct_patient_count(logic, date_range, patient_type)
            return response.total
        response = cohort_api.selector_patient_count(logic, patient_type, date_range)
        return response.total_count

    logic_desc = describe_logic_tree(logic)
    cohort_label = logic_desc.split(":")[-1].strip() if ":" in logic_desc else "Query"
    count_type_label = "distinct" if count_type == CountType.DISTINCT else "approx"

    console.print("")
    console.print("\u2550" * 79)
    title = f"COHORT SUMMARY: {cohort_label}"
    console.print(f"{title:^79}")
    console.print("\u2550" * 79)
    console.print("")
    console.print(f"  Date Range: {date_range.from_date.isoformat()} to {date_range.to_date.isoformat()}")
    console.print(f"  [dim]Loading breakdown data...[/dim]")

    section_definitions = [
        ("total", ["total_adv"]),
        ("yearly", ["yearly"]),
        ("quarterly", ["quarterly"]),
        ("gender", ["sex_M", "sex_F"]),
        ("age", [f"age_{label}" for label, _, _ in AGE_BUCKETS]),
        ("states", [f"state_{state}" for state in all_states]),
        ("race", [f"race_{race}" for race in RACE_ETHNICITY]),
        ("mortality", ["deceased"]),
    ]

    sections_printed: set[str] = set()

    def query_by_sex(sex: str) -> int:
        criteria = {"table": "demo", "filter": {"col": "patient_sex", "op": "eq", "value": sex}}
        return cohort_api.distinct_patient_count(logic, date_range, patient_criteria=criteria).total

    def query_by_state(state: str) -> int:
        criteria = {"table": "geo", "timing": "eventdate", "filter": {"col": "patient_state", "op": "eq", "value": state}}
        return cohort_api.distinct_patient_count(logic, date_range, patient_criteria=criteria).total

    def query_by_age(age_from: int, age_to: int) -> int:
        criteria = {"table": "demo", "timing": "eventdate", "filter": {"col": "patient_age", "op": "in_ranges", "value": [{"from": age_from, "to": age_to}]}}
        return cohort_api.distinct_patient_count(logic, date_range, patient_criteria=criteria).total

    def query_by_race(race: str) -> int:
        criteria = {"table": "re", "filter": {"col": "patient_race_ethnicity", "op": "eq", "value": race}}
        return cohort_api.distinct_patient_count(logic, date_range, patient_criteria=criteria).total

    def query_deceased() -> int:
        criteria = {
            "table": "mortality",
            "timing": "first",
            "filter": {
                "op": "and",
                "conditions": [
                    {"col": "patient_death_date", "op": "in_date_ranges", "value": [{"from": date_range.from_date.isoformat(), "to": date_range.to_date.isoformat()}]},
                    {"col": "has_multiple_death_dates", "op": "not_eq", "value": True},
                ],
            },
        }
        return cohort_api.distinct_patient_count(logic, date_range, patient_criteria=criteria).total

    def query_by_time(granularity: str) -> dict[str, int]:
        return cohort_api.distinct_patient_count(logic, date_range, granularity=granularity).counts

    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        futures: dict[concurrent.futures.Future, str] = {}

        f = executor.submit(get_count_for_type, "adv")
        futures[f] = "total_adv"
        query_count += 1

        f = executor.submit(query_by_time, "year")
        futures[f] = "yearly"
        query_count += 1

        f = executor.submit(query_by_time, "quarter")
        futures[f] = "quarterly"
        query_count += 1

        for sex in ["M", "F"]:
            f = executor.submit(query_by_sex, sex)
            futures[f] = f"sex_{sex}"
            query_count += 1

        for label, age_from, age_to in AGE_BUCKETS:
            f = executor.submit(query_by_age, age_from, age_to)
            futures[f] = f"age_{label}"
            query_count += 1

        for state in all_states:
            f = executor.submit(query_by_state, state)
            futures[f] = f"state_{state}"
            query_count += 1

        for race in RACE_ETHNICITY:
            f = executor.submit(query_by_race, race)
            futures[f] = f"race_{race}"
            query_count += 1

        f = executor.submit(query_deceased)
        futures[f] = "deceased"
        query_count += 1

        try:
            for future in concurrent.futures.as_completed(futures, timeout=300):
                key = futures[future]
                try:
                    results[key] = future.result()
                except Exception as e:
                    errors[key] = e
                    logger.warning(f"Query {key} failed: {e}")

                for section_name, query_keys in section_definitions:
                    if section_name in sections_printed:
                        continue

                    if all((qk in results) or (qk in errors) for qk in query_keys):
                        _print_section_streaming(section_name, results, errors, console, all_states, count_type_label)
                        sections_printed.add(section_name)
        except concurrent.futures.TimeoutError:
            logger.warning("Timeout while waiting for cohort exploration queries; some queries did not complete within 300s.")
            for future, key in futures.items():
                if (key not in results) and (key not in errors):
                    future.cancel()
                    errors[key] = TimeoutError("Query timed out after 300 seconds.")

            for section_name, query_keys in section_definitions:
                if section_name in sections_printed:
                    continue
                if all((qk in results) or (qk in errors) for qk in query_keys):
                    _print_section_streaming(section_name, results, errors, console, all_states, count_type_label)
                    sections_printed.add(section_name)
    elapsed = (datetime.now() - start_time).total_seconds()
    console.print("")
    console.print("\u2550" * 79)
    console.print(f"  Query completed in {elapsed:.1f}s ({query_count} queries) | {datetime.now().strftime('%Y-%m-%d')}")
    console.print("\u2550" * 79)


def _print_section_streaming(
    section_name: str,
    results: dict,
    errors: dict[str, Exception],
    console: Console,
    all_states: list[str],
    count_type_label: str,
) -> None:
    """Print a single section of the comprehensive summary."""
    total_adv = results.get("total_adv")
    total_adv_value = total_adv if total_adv is not None else 0
    total_adv_error = errors.get("total_adv")

    def format_percent_or_na(value: int) -> str:
        if total_adv is None or total_adv_error:
            return "n/a"
        return format_percent(value, total_adv)

    if section_name == "total":
        if total_adv_error:
            console.print(f"  Total Patients ({count_type_label}): (failed: {total_adv_error})")
        else:
            console.print(f"  Total Patients ({count_type_label}): {format_number(total_adv_value)}")

    elif section_name == "yearly":
        yearly_error = errors.get("yearly")
        yearly = results.get("yearly") or {}
        console.print("")
        console.print("\u2500" * 79)
        console.print("  TREND BY YEAR")
        console.print("\u2500" * 79)
        if yearly_error:
            console.print(f"  (Failed to load: {yearly_error})")
        elif yearly:
            console.print("")
            max_val = max(yearly.values()) if yearly.values() else 1
            for year in sorted(yearly.keys()):
                count = yearly[year]
                bar = make_bar(count, max_val, 45)
                console.print(f"  {year} \u2502{bar}  {format_number(count):>12}")
                console.print("")
        else:
            console.print("  (No data)")

    elif section_name == "quarterly":
        quarterly_error = errors.get("quarterly")
        quarterly = results.get("quarterly") or {}
        console.print("")
        console.print("\u2500" * 79)
        console.print("  TREND BY QUARTER")
        console.print("\u2500" * 79)
        if quarterly_error:
            console.print(f"  (Failed to load: {quarterly_error})")
        elif quarterly:
            console.print("")
            max_val = max(quarterly.values()) if quarterly.values() else 1
            for q in sorted(quarterly.keys()):
                count = quarterly[q]
                bar = make_bar(count, max_val, 45)
                console.print(f"  {q} \u2502{bar}  {format_number(count):>12}")
                console.print("")
        else:
            console.print("  (No data)")

    elif section_name == "gender":
        male_error = errors.get("sex_M")
        female_error = errors.get("sex_F")
        male = results.get("sex_M")
        female = results.get("sex_F")
        console.print("")
        console.print("\u2500" * 79)
        console.print("  GENDER")
        console.print("\u2500" * 79)
        if male is None and female is None and (male_error or female_error):
            if male_error:
                console.print(f"  Male \u2502 (failed: {male_error})")
            if female_error:
                console.print(f"  Female \u2502 (failed: {female_error})")
        elif male is not None or female is not None:
            console.print("")
            max_val = max(v for v in [male, female] if v is not None)
            if male is not None:
                console.print(f"  {'Male':>6} \u2502{make_bar(male, max_val, 40)}  {format_number(male):>12} ({format_percent_or_na(male)})")
            elif male_error:
                console.print(f"  {'Male':>6} \u2502  (failed: {male_error})")
            console.print("")
            if female is not None:
                console.print(f"  {'Female':>6} \u2502{make_bar(female, max_val, 40)}  {format_number(female):>12} ({format_percent_or_na(female)})")
            elif female_error:
                console.print(f"  {'Female':>6} \u2502  (failed: {female_error})")
            console.print("")
        else:
            console.print("  (No data)")

    elif section_name == "age":
        age_data = {label: results.get(f"age_{label}") for label, _, _ in AGE_BUCKETS}
        age_values = [v for v in age_data.values() if v is not None]
        console.print("")
        console.print("\u2500" * 79)
        console.print("  AGE DISTRIBUTION")
        console.print("\u2500" * 79)
        if age_values:
            console.print("")
            max_val = max(age_values)
            for label, _, _ in AGE_BUCKETS:
                count = age_data.get(label)
                if count is not None:
                    bar = make_bar(count, max_val, 40)
                    console.print(f"  {label:>5} \u2502{bar}  {format_number(count):>12} ({format_percent_or_na(count)})")
                else:
                    error = errors.get(f"age_{label}")
                    if error:
                        console.print(f"  {label:>5} \u2502  (failed: {error})")
                    else:
                        console.print(f"  {label:>5} \u2502  (no data)")
                console.print("")
        else:
            console.print("  (No data)")

    elif section_name == "states":
        has_additional_states = len(all_states) > len(TOP_10_STATES)
        console.print("")
        console.print("\u2500" * 79)
        if has_additional_states:
            console.print("  STATES (top 10 by US population + custom)")
        else:
            console.print("  TOP 10 STATES (by US population)")
        console.print("\u2500" * 79)
        state_data = [(state, results.get(f"state_{state}")) for state in all_states]
        state_data = [(s, c) for s, c in state_data if c is not None]
        state_data.sort(key=lambda x: x[1], reverse=True)
        if state_data:
            console.print("")
            max_val = max(c for _, c in state_data)
            for state, count in state_data:
                bar = make_bar(count, max_val, 40)
                console.print(f"  {state:>2} \u2502{bar}  {format_number(count):>12} ({format_percent_or_na(count)})")
                console.print("")
            failed_states = [state for state in all_states if f"state_{state}" in errors]
            if failed_states:
                console.print(f"  (Failed: {', '.join(failed_states)})")
        else:
            failed_states = [state for state in all_states if f"state_{state}" in errors]
            if failed_states:
                console.print(f"  (Failed: {', '.join(failed_states)})")
            else:
                console.print("  (No data)")

    elif section_name == "race":
        console.print("")
        console.print("\u2500" * 79)
        console.print("  RACE / ETHNICITY")
        console.print("\u2500" * 79)
        race_data = [(race, results.get(f"race_{race}")) for race in RACE_ETHNICITY]
        race_data = [(r, c) for r, c in race_data if c is not None]
        race_data.sort(key=lambda x: x[1], reverse=True)
        if race_data:
            console.print("")
            max_val = max(c for _, c in race_data)
            for race, count in race_data:
                bar = make_bar(count, max_val, 35)
                label = race[:25] + "..." if len(race) > 25 else race
                console.print(f"  {label:>28} \u2502{bar}  {format_number(count):>12} ({format_percent_or_na(count)})")
                console.print("")
            failed_races = [race for race in RACE_ETHNICITY if f"race_{race}" in errors]
            if failed_races:
                console.print(f"  (Failed: {', '.join(failed_races)})")
        else:
            failed_races = [race for race in RACE_ETHNICITY if f"race_{race}" in errors]
            if failed_races:
                console.print(f"  (Failed: {', '.join(failed_races)})")
            else:
                console.print("  (No data)")

    elif section_name == "mortality":
        deceased_error = errors.get("deceased")
        deceased = results.get("deceased")
        console.print("")
        console.print("\u2500" * 79)
        console.print("  MORTALITY")
        console.print("\u2500" * 79)
        if deceased_error:
            console.print(f"  (Failed to load: {deceased_error})")
        elif deceased is not None:
            console.print("")
            if total_adv is None or total_adv_error:
                console.print(f"  {'Living':>8} \u2502  (total count unavailable)")
                console.print("")
                console.print(f"  {'Deceased':>8} \u2502  {format_number(deceased):>12} (n/a)")
                console.print("")
            else:
                living = total_adv_value - deceased
                max_val = max(living, deceased) if living > 0 or deceased > 0 else 1
                console.print(f"  {'Living':>8} \u2502{make_bar(living, max_val, 40)}  {format_number(living):>12} ({format_percent(living, total_adv)})")
                console.print("")
                console.print(f"  {'Deceased':>8} \u2502{make_bar(deceased, max_val, 40)}  {format_number(deceased):>12} ({format_percent(deceased, total_adv)})")
                console.print("")
        else:
            console.print("  (No data)")


# ---------------------------------------------------------------------------
# Query execution
# ---------------------------------------------------------------------------


def execute_cohort_query(
    cohort_api: CohortApi,
    params: dict,
    console: Console,
) -> dict:
    """Execute the cohort count query against the Cohort API."""
    from komodo.exceptions import ApiException

    logic = params.get("logic")
    if not logic:
        console.print(Text("\u274c No cohort logic provided.", style="bold red"))
        raise typer.Exit(code=1)

    count_type = params.get("count_type", CountType.APPROXIMATE)

    date_range = None
    if params.get("date_from") and params.get("date_to"):
        date_range = DateRange.from_strings(params["date_from"], params["date_to"])
    elif params.get("date_from"):
        date_range = DateRange.from_strings(params["date_from"])

    if count_type == CountType.DISTINCT and not date_range:
        console.print(Text("\u274c Date range is required for distinct patient counts.", style="bold red"))
        raise typer.Exit(code=1)

    try:
        count, response = cohort_api.get_total_count(
            logic=logic,
            count_type=count_type,
            date_range=date_range,
        )
        result = {
            "count_type": count_type,
            "totalCount": count,
            "warnings": getattr(response, "warnings", []),
        }
        return result
    except ApiException as e:
        if e.status == 401:
            console.print(Text("\u274c Authentication failed. Please run 'komodo login' first.", style="bold red"))
        elif e.status == 403:
            console.print(Text("\u274c Access denied. You don't have permission to query cohorts.", style="bold red"))
        else:
            console.print(Text(f"\u274c API error ({e.status}): {e.reason}", style="bold red"))
        raise typer.Exit(code=1)
    except Exception as e:
        console.print(Text(f"\u274c Error querying Cohort API: {e}", style="bold red"))
        raise typer.Exit(code=1)


def display_results(params: dict, result: dict, console: Console) -> None:
    """Display the cohort query results in a formatted panel."""
    count = result.get("totalCount", result.get("count", result.get("total", "N/A")))
    warnings = result.get("warnings", [])
    count_type = result.get("count_type", params.get("count_type", CountType.APPROXIMATE))

    logic = params.get("logic", {})
    criteria_text = describe_logic_tree(logic)

    if params.get("date_from"):
        date_range = params["date_from"]
        if params.get("date_to"):
            date_range += f" to {params['date_to']}"
        criteria_text += f"\n  \u2022 Date Range: [cyan]{date_range}[/cyan]"

    if isinstance(count, (int, float)):
        formatted_count = f"{int(count):,}"
    else:
        formatted_count = str(count)

    count_type_label = "Distinct" if count_type == CountType.DISTINCT else "Approximate"
    count_emoji = "\U0001f3af" if count_type == CountType.DISTINCT else "\u26a1"

    result_text = f"[bold green]{count_emoji} {count_type_label} Patient Count: {formatted_count}[/bold green]\n\n" f"[bold]Query Criteria:[/bold]\n{criteria_text}"

    if warnings:
        warning_text = "\n\n[bold yellow]\u26a0\ufe0f  Warnings:[/bold yellow]\n"
        for warning in warnings:
            warning_text += f"  \u2022 {warning}\n"
        result_text += warning_text

    console.print(
        Panel(
            result_text,
            title="\U0001f4ca Cohort Exploration Results",
            border_style="green",
        )
    )


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


@resolve_environment
def cohort_exploration(environment: str = ENVIRONMENT_OPTION):
    """
    [Beta] Explore patient cohorts interactively using the Cohort API.

    This is a beta feature that may change or be removed without notice.
    Access is limited to select users during the beta period.

    This command provides a streamlined interface for rapid validation of clinical
    hypotheses via interactive prompts. Generates a comprehensive cohort summary
    with breakdowns by time trends, gender, age, state, race/ethnicity, and mortality.

    Supported parameters:
    - Diagnosis Code (ICD-10 format, e.g., C34%, C50.1)
    - Procedure Code (CPT/HCPCS format, e.g., 96413, J9271)
    - Drug Reference (NDC Code, e.g., 12345-6789-01)
    - Time Duration (date range in mm/dd/yy format)
    """
    from komodo.auth.Session import Session
    from komodo.client import Client
    from komodo.exceptions import ApiException

    console = Console()

    session = Session(environment=environment)
    credentials = session.load_credentials()

    if not credentials or not credentials.access_token:
        console.print(Text("\u274c No access token found. Please run 'komodo login' first.", style="bold red"))
        raise typer.Exit(code=1)

    client = Client(auth_session=session)
    cohort_api = CohortApi(api_client=client.session)

    try:
        params = prompt_for_input(console)
    except KeyboardInterrupt:
        console.print(Text("\n\U0001f44b Cohort exploration cancelled.", style="yellow"))
        raise typer.Exit(code=0)

    logic = params.get("logic")
    if not logic:
        console.print(Text("\u274c No cohort logic provided.", style="bold red"))
        raise typer.Exit(code=1)

    output_mode = params.get("output_mode", "comprehensive")
    count_type = params.get("count_type", CountType.APPROXIMATE)

    if (output_mode == "comprehensive" or count_type == CountType.DISTINCT) and not params.get("date_from"):
        console.print(Text("\u274c Date range is required for comprehensive summary or distinct counts.", style="bold red"))
        raise typer.Exit(code=1)

    date_range = None
    if params.get("date_from"):
        date_range = DateRange.from_strings(
            params["date_from"],
            params.get("date_to"),
        )

    logic_desc = describe_logic_tree(logic)
    cohort_label = logic_desc.split(":")[-1].strip() if ":" in logic_desc else "cohort"

    if output_mode == "quick":
        console.print(f"\n  Fetching patient count for {cohort_label}...")

        try:
            result = execute_cohort_query(
                cohort_api=cohort_api,
                params=params,
                console=console,
            )
            display_results(params, result, console)
        except typer.Exit:
            raise
        except Exception as e:
            console.print(Text(f"\u274c Unexpected error: {e}", style="bold red"))
            logger.exception("Cohort exploration error")
            raise typer.Exit(code=1)
    else:
        additional_states = params.get("additional_states", [])
        try:
            run_comprehensive_summary_streaming(cohort_api, logic, date_range, console, additional_states, count_type)
        except ApiException as e:
            if e.status == 401:
                console.print(Text("\u274c Authentication failed. Please run 'komodo login' first.", style="bold red"))
            elif e.status == 403:
                console.print(Text("\u274c Access denied. You don't have permission to query cohorts.", style="bold red"))
            else:
                console.print(Text(f"\u274c API error ({e.status}): {e.reason}", style="bold red"))
            raise typer.Exit(code=1)
        except Exception as e:
            console.print(Text(f"\u274c Unexpected error: {e}", style="bold red"))
            logger.exception("Cohort exploration error")
            raise typer.Exit(code=1)
