"""
This module handles LLM querying to annotate cell types in the context of all other cells, based on marker genes alone.
"""

from anndict.utils import enforce_semantic_list
from anndict.llm.llm_call import call_llm
from .base import ai_annotate

def ai_cell_type(gene_list, tissue=None):
    """
    Returns the cell type based on a list of marker genes as determined by AI.
    Args:
    gene_list (list of str): The list of genes to be described.
    tissue (str, optional): The tissue of origin to provide context for the AI.
    Returns:
    str: The cell type label generated by AI
    """
    #enforce that labels are semantic
    enforce_semantic_list(gene_list)

    # Prepare the prompt
    if len(gene_list) == 1:
        base_prompt = f"In a few words and without restating any part of the question, describe the single most likely cell type represented by the marker gene: {gene_list[0]}"
    else:
        genes_str = "    ".join(gene_list)
        base_prompt = f"In a few words and without restating any part of the question, describe the single most likely cell type represented by the marker genes: {genes_str}"

    # Add tissue information if provided
    if tissue:
        base_prompt += f" Consider that these cells are from {tissue} tissue."

    # Prepare the messages for the Chat Completions API
    messages = [
        {"role": "system", "content": "You are a terse molecular biologist."},
        {"role": "user", "content": base_prompt}
    ]

    # Call the LLM using the call_llm function
    annotation = call_llm(
        messages=messages,
        max_tokens=100,
        temperature=0
    )

    return annotation

def ai_annotate_cell_type(adata, groupby, n_top_genes, label_column='ai_cell_type', tissue_of_origin_col=None):
    """
    Annotate cell types based on the top marker genes for each cluster.

    This function performs differential expression analysis to identify marker genes for each cluster
    and applies ai_annotate_cell_type to determine the cell type for each cluster based on the top 
    marker genes. The results are added to the AnnData object and returned as a DataFrame.

    Parameters:
    adata : AnnData
    groupby : str Column in adata.obs to group by for differential expression analysis.
    n_top_genes : int The number of top marker genes to consider for each cluster.
    label_column : str, optional (default: 'ai_cell_type') The name of the new column in adata.obs where the cell type annotations will be stored.

    Returns:
    pd.DataFrame A DataFrame with a column for the top marker genes for each cluster.
    """
    return ai_annotate(func=ai_cell_type, adata=adata, groupby=groupby, n_top_genes=n_top_genes, label_column=label_column, tissue_of_origin_col=tissue_of_origin_col)
