"""Stage 3: Thematic Reason Clustering."""
from __future__ import annotations
from .sysmsg import make_clustering_system_message
from .schemas import build_cluster_schema
from .constants import CLUSTERING_OUTPUT_TOKENS, MAX_NUM_CLUSTER
from crystallise.llm.client import batch_chat_completions
from crystallise.common.json_utils import parse_json_safe as _parse_json_safe


def make_clustering_input(reasoning_list: list[str], inc_excl_dict: dict, decision_type: str):
    output_schema = build_cluster_schema(list(inc_excl_dict.keys()))
    system_message = make_clustering_system_message(inc_excl_dict.keys(), decision_type, max_clusters=MAX_NUM_CLUSTER)
    formatted_text = "\n\n>\n\n".join(reasoning_list)
    return system_message, formatted_text, output_schema


def get_reasoning_clusters(
    reasoning_list: list[str],
    model_name: str,
    inc_excl_dict: dict,
    decision_type: str,
    cancel_event=None,
) -> list[dict]:
    """Cluster reasoning texts into themes. Returns list of cluster dicts."""
    if cancel_event is not None and cancel_event.is_set():
        return []
    system_message, formatted_text, output_schema = make_clustering_input(
        reasoning_list, inc_excl_dict, decision_type
    )
    response = batch_chat_completions(
        system_messages=[system_message],
        prompts=[formatted_text],
        output_schemas=[output_schema],
        max_completion_tokens=CLUSTERING_OUTPUT_TOKENS,
        model=model_name,
    )
    response_json = _parse_json_safe(response[0])
    cluster_list = response_json.get("cluster_list") if isinstance(response_json, dict) else None
    if cluster_list is None:
        raise RuntimeError(f"Clustering failed: cluster_list is None for {decision_type}")
    return cluster_list