import random
from typing import List, Dict, Union
from smolagents.tools import tool
import pandas as pd

@tool
def generate_data_table(
    topic: str, 
    columns: List[str], 
    num_rows: int, 
    data_ranges: Dict[str, List[Union[int, str]]]
) -> Dict:
    """
    Gera uma tabela de dados fictícios sobre um determinado tópico.
    Útil para criar dados para questões que envolvem tabelas de dupla entrada e gráficos (EF02MA22, EF03MA26).

    Args:
        topic (str): O tópico geral da tabela (ex: "Vendas de Frutas", "Notas de Alunos").
        columns (List[str]): Uma lista com os nomes das colunas (ex: ["Mês", "Maçãs", "Bananas"]).
        num_rows (int): O número de linhas (registros) a serem gerados.
        data_ranges (Dict): Um dicionário onde cada chave é um nome de coluna e o valor é uma lista 
                            de possíveis valores ou um intervalo [min, max] para números.
                            Exemplo: {"Mês": ["Jan", "Fev", "Mar"], "Maçãs": [100, 500]}

    Returns:
        Um dicionário contendo a tabela como uma lista de dicionários e também em formato Markdown.
    """
    table_data = []
    for _ in range(num_rows):
        row = {}
        for col in columns:
            if col not in data_ranges:
                row[col] = None
                continue
            
            range_info = data_ranges[col]
            if isinstance(range_info[0], str):
                # Se for uma lista de strings, escolhe uma aleatoriamente
                row[col] = random.choice(range_info)
            elif isinstance(range_info[0], int) and len(range_info) == 2:
                # Se for uma lista de 2 números, gera um inteiro aleatório no intervalo
                row[col] = random.randint(range_info[0], range_info[1])
        table_data.append(row)
        
    # Converte para um DataFrame do Pandas para fácil formatação
    df = pd.DataFrame(table_data)
    
    # Formata a tabela em Markdown para fácil visualização
    markdown_table = df.to_markdown(index=False)
    
    return {
        "topic": topic,
        "table_as_list_of_dicts": table_data,
        "table_as_markdown": markdown_table
    }