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

from pandas import DataFrame
from anndata import AnnData

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: list[str],
    tissue: str = None
) -> str:
    """
    Returns the cell type based on a list of marker genes as determined by LLM.

    Parameters
    -------------
    gene_list
        The list of genes to be described.

    tissue
        The tissue of origin to provide context for the LLM.

    Returns
    --------
    The cell type label generated by an LLM.
    """
    if not gene_list:
        raise ValueError("Empty gene list passed to ai_cell_type")

    #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: AnnData,
    groupby: str,
    n_top_genes: int,
    new_label_column: str = 'ai_cell_type',
    tissue_of_origin_col: str = None
) -> DataFrame:
    """
    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
        Column in adata.obs to group by for differential expression analysis.
    groupby
        Column in adata.obs to group by for differential expression analysis.
    n_top_genes
        The number of top marker genes to consider for each cluster.
    new_label_column
        The name of the new column in adata.obs where the cell type annotations will be stored.
        Default: 'ai_cell_type'
    tissue_of_origin_col
        Missing description for this parameter.
        Default: None

    Returns
    -------
    A pd.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, new_label_column=new_label_column, tissue_of_origin_col=tissue_of_origin_col)
