import gradio as gr
import tiktoken
import pandas as pd
import os
from datasets import load_dataset

# Carica il dataset di prezzi
def load_llm_prices():
    # try:
    #     # Ottieni il token Hugging Face dalle variabili d'ambiente/secrets
    #     hf_token = os.environ.get("HUGGINGFACE_TOKEN")
        
    #     # Nome del dataset (repository/dataset_name)
    #     dataset_name = os.environ.get("DATASET_NAME")
        
    #     # Carica il dataset usando la libreria datasets
    #     dataset = load_dataset(dataset_name, token=hf_token)
        
    #     # Converti il dataset in un DataFrame
    #     df = pd.DataFrame(dataset["train"])
        
    #     print(f"Dataset caricato con successo: {len(df)} record")
    #     print(f"Colonne disponibili: {df.columns.tolist()}")
    #     print(f"Provider disponibili: {df['provider'].unique().tolist()}")
        
    #     return df
    # except Exception as e:
    #     print(f"Errore nel caricamento del dataset: {e}")
    #     # Creiamo un dataset di esempio in caso di errore
    #     return pd.DataFrame({
    #         'provider': [
    #             'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI',
    #             'Anthropic', 'Anthropic', 'Anthropic'
    #         ],
    #         'model_id': [
    #             'gpt-5', 'gpt-5-mini', 'gpt-5-nano', 'gpt-5-chat-latest', 'gpt-4.1', 'gpt-4.1-mini', 
    #             'gpt-4.1-nano', 'gpt-4o', 'gpt-4o-2024-05-13', 'gpt-4o-mini',
    #             'claude-4-opus', 'claude-4-sonnet', 'claude-3.5-haiku'
    #         ],
    #         'price_per_input_token': [
    #             0.00000125, 0.00000025, 0.00000005, 0.00000125, 0.000002, 0.0000004, 
    #             0.0000001, 0.0000025, 0.000005, 0.00000015,
    #             0.000015, 0.000003, 0.0000008
    #         ],
    #         'price_per_output_token': [
    #             0.0000125, 0.0000025, 0.0000005, 0.0000125, 0.00000050, 0.0000001, 
    #             0.0000025, 0.0000125, 0.000015, 0.0000075,
    #             0.000075, 0.000015, 0.000004
    #         ]
    #     })
    return pd.DataFrame({
            'provider': [
                'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI', 'OpenAI'            ],
            'model_id': [
                'gpt-5', 'gpt-5-mini', 'gpt-5-nano', 'gpt-5-chat-latest', 'gpt-4.1', 'gpt-4.1-mini', 
                'gpt-4.1-nano', 'gpt-4o', 'gpt-4o-2024-05-13', 'gpt-4o-mini',
            ],
            'price_per_input_token': [
                0.00000125, 0.00000025, 0.00000005, 0.00000125, 0.000002, 0.0000004, 
                0.0000001, 0.0000025, 0.000005, 0.00000015            ],
            'price_per_output_token': [
                0.0000125, 0.0000025, 0.0000005, 0.0000125, 0.00000050, 0.0000001, 
                0.0000025, 0.0000125, 0.000015, 0.0000075            ]
        })

# Carica i prezzi degli LLM
llm_prices_df = load_llm_prices()
unique_providers = sorted(llm_prices_df['provider'].unique().tolist())
print(f"Provider caricati: {unique_providers}")

# Funzione per ottenere i modelli per un provider
def get_models(provider):
    print(f"Cercando modelli per: {provider}")
    if not provider or provider == "Nessuno":
        return ["Nessuno"]
    
    filtered_models = llm_prices_df[llm_prices_df['provider'] == provider]['model_id'].tolist()
    print(f"Modelli trovati per {provider}: {filtered_models}")
    
    return filtered_models if filtered_models else ["Nessuno"]

# Ottieni i modelli iniziali per il primo provider
initial_provider = unique_providers[0] if unique_providers else "Nessuno"
initial_models = get_models(initial_provider)
initial_model = initial_models[0] if initial_models else "Nessuno"
print(f"Provider iniziale: {initial_provider}, Modello iniziale: {initial_model}")

# Calcolo dei token e prezzi
def count_tokens_and_price(input_text, output_text, provider1, model_id1, provider2, model_id2):
    print(f"Calcolo costi con: Provider1={provider1}, Model1={model_id1}, Provider2={provider2}, Model2={model_id2}")
    
    encoding = tiktoken.get_encoding("cl100k_base")
    
    # Conta i token
    input_tokens = len(encoding.encode(input_text))
    output_tokens = len(encoding.encode(output_text))
    total_tokens = input_tokens + output_tokens
    
    # Calcola il confronto
    ratio = output_tokens/input_tokens if input_tokens > 0 else 0
    
    # Funzione per calcolare i costi di un singolo modello
    def calculate_model_costs(provider, model_id):
        if provider == "Nessuno" or model_id == "Nessuno":
            return None
            
        model_data = llm_prices_df[(llm_prices_df['provider'] == provider) & 
                                  (llm_prices_df['model_id'] == model_id)]
        
        if not model_data.empty:
            input_price_per_token = model_data['price_per_input_token'].values[0]
            output_price_per_token = model_data['price_per_output_token'].values[0]
            
            # Calcola i costi
            input_cost = input_tokens * input_price_per_token
            output_cost = output_tokens * output_price_per_token
            total_cost = input_cost + output_cost
            
            return {
                'provider': provider,
                'model': model_id,
                'input_price': input_price_per_token,
                'output_price': output_price_per_token,
                'input_cost': input_cost,
                'output_cost': output_cost,
                'total_cost': total_cost
            }
        else:
            return None
    
    # Calcola i costi per entrambi i modelli
    model1_costs = calculate_model_costs(provider1, model_id1)
    model2_costs = calculate_model_costs(provider2, model_id2)
    
    # Informazioni sui token
    token_info = f"""📊 Analisi dei Token:
- Input tokens: {input_tokens}
- Output tokens: {output_tokens}
- Total tokens: {total_tokens}
- Ratio input/output: {ratio:.2f}x"""
    
    # Prezzi per token in formato tabella
    if model1_costs:
        model1_info = f"| **Modello 1** | **{model1_costs['provider']} - {model1_costs['model']}** |\n|-----------|------------:|\n| Prezzo per token di input | ${model1_costs['input_price']:.6f} |\n| Prezzo per token di output | ${model1_costs['output_price']:.6f} |"
        
        if model2_costs:
            model2_info = f"\n\n| **Modello 2** | **{model2_costs['provider']} - {model2_costs['model']}** |\n|-----------|------------:|\n| Prezzo per token di input | ${model2_costs['input_price']:.6f} |\n| Prezzo per token di output | ${model2_costs['output_price']:.6f} |"
            price_per_token_info = model1_info + model2_info
        else:
            price_per_token_info = model1_info
    else:
        price_per_token_info = "⚠️ Modello 1 non trovato nel dataset"
    
    # Costruisci la tabella di confronto
    if model2_costs and model1_costs:
        # Confronto tra due modelli
        diff_input = abs(model1_costs['input_cost'] - model2_costs['input_cost'])
        diff_output = abs(model1_costs['output_cost'] - model2_costs['output_cost'])
        diff_total = abs(model1_costs['total_cost'] - model2_costs['total_cost'])
        
        percentage_diff = (100 * diff_total / max(model1_costs['total_cost'], model2_costs['total_cost']))
        
        cheaper_model = model1_costs['model'] if model1_costs['total_cost'] <= model2_costs['total_cost'] else model2_costs['model']
        
        token_cost_table = f"""### 📊 Tabella Comparativa

| Tipo      | Numero    | {model1_costs['provider']} - {model1_costs['model']} | {model2_costs['provider']} - {model2_costs['model']} | Differenza    |
|:----------|----------:|------------:|------------:|------------:|
| **Input**     | {input_tokens:,} | ${model1_costs['input_cost']:.6f} | ${model2_costs['input_cost']:.6f} | ${diff_input:.6f} |
| **Output**    | {output_tokens:,} | ${model1_costs['output_cost']:.6f} | ${model2_costs['output_cost']:.6f} | ${diff_output:.6f} |
| **Totale**    | {total_tokens:,} | ${model1_costs['total_cost']:.6f} | ${model2_costs['total_cost']:.6f} | ${diff_total:.6f} |

### 💰 Analisi di Risparmio
- **Risparmio totale:** ${diff_total:.6f}
- **Risparmio percentuale:** {percentage_diff:.2f}%
- **Modello più economico:** {cheaper_model}
"""
    elif model1_costs:
        # Solo il primo modello
        token_cost_table = f"""### 📊 Tabella dei Costi

| Tipo      | Numero    | {model1_costs['provider']} - {model1_costs['model']} |
|:----------|----------:|------------:|
| **Input**     | {input_tokens:,} | ${model1_costs['input_cost']:.6f} |
| **Output**    | {output_tokens:,} | ${model1_costs['output_cost']:.6f} |
| **Totale**    | {total_tokens:,} | ${model1_costs['total_cost']:.6f} |"""
    else:
        token_cost_table = "⚠️ Dati non disponibili"
    
    # Ritorna le stringhe formattate
    return str(total_tokens), token_info, price_per_token_info, token_cost_table

# Generazione esempio di testo con token specifici
def show_text_length(ntokens, provider1, model_id1, provider2, model_id2):
    lorem_ipsum = """Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."""
    encoding = tiktoken.get_encoding("cl100k_base")
    tokens = encoding.encode(lorem_ipsum)
    token_count = len(tokens)
    
    # Aggiusta il testo per raggiungere il numero desiderato di token
    while token_count < int(ntokens):
        lorem_ipsum += " " + lorem_ipsum
        tokens = encoding.encode(lorem_ipsum)
        token_count = len(tokens)
    
    # Taglia il testo se necessario per ottenere il numero esatto di token
    while token_count > int(ntokens):
        lorem_ipsum = lorem_ipsum[:-1]
        tokens = encoding.encode(lorem_ipsum)
        token_count = len(tokens)
    
    # Funzione per calcolare i costi di un singolo modello
    def calculate_model_costs(provider, model_id):
        if provider == "Nessuno" or model_id == "Nessuno":
            return None
            
        model_data = llm_prices_df[(llm_prices_df['provider'] == provider) & 
                                  (llm_prices_df['model_id'] == model_id)]
        
        if not model_data.empty:
            input_price_per_token = model_data['price_per_input_token'].values[0]
            output_price_per_token = model_data['price_per_output_token'].values[0]
            
            # Calcola i costi
            input_cost = token_count * input_price_per_token
            output_cost = token_count * output_price_per_token
            
            return {
                'provider': provider,
                'model': model_id,
                'input_price': input_price_per_token,
                'output_price': output_price_per_token,
                'input_cost': input_cost,
                'output_cost': output_cost
            }
        else:
            return None
    
    # Calcola i costi per entrambi i modelli
    model1_costs = calculate_model_costs(provider1, model_id1)
    model2_costs = calculate_model_costs(provider2, model_id2)
    
    # Costruisci la tabella di confronto
    if model1_costs:
        if model2_costs:
            # Confronto tra due modelli
            diff_input = abs(model1_costs['input_cost'] - model2_costs['input_cost'])
            diff_output = abs(model1_costs['output_cost'] - model2_costs['output_cost'])
            
            cheaper_input = model1_costs['model'] if model1_costs['input_cost'] <= model2_costs['input_cost'] else model2_costs['model']
            cheaper_output = model1_costs['model'] if model1_costs['output_cost'] <= model2_costs['output_cost'] else model2_costs['model']
            
            token_cost_table = f"""### 📊 Tabella Comparativa per {token_count:,} Token

| Scenario | {model1_costs['provider']} - {model1_costs['model']} | {model2_costs['provider']} - {model2_costs['model']} | Differenza |
|:---------|------------:|------------:|------------:|
| **Se Input**    | ${model1_costs['input_cost']:.6f} | ${model2_costs['input_cost']:.6f} | ${diff_input:.6f} |
| **Se Output**   | ${model1_costs['output_cost']:.6f} | ${model2_costs['output_cost']:.6f} | ${diff_output:.6f} |

### 💰 Analisi di Risparmio
- **Risparmio (se input):** ${diff_input:.6f}
- **Risparmio (se output):** ${diff_output:.6f}
- **Modello più economico (input):** {cheaper_input}
- **Modello più economico (output):** {cheaper_output}
"""
        else:
            # Solo il primo modello
            token_cost_table = f"""### 📊 Costi Stimati per {token_count:,} Token

| Scenario | {model1_costs['provider']} - {model1_costs['model']} |
|:---------|------------:|
| **Se Input** | ${model1_costs['input_cost']:.6f} |
| **Se Output**| ${model1_costs['output_cost']:.6f} |"""
    else:
        token_cost_table = "⚠️ Seleziona un provider e un modello per vedere i costi"
    
    return lorem_ipsum, str(token_count), token_cost_table

# Aggiorna modelli e seleziona il primo
def update_models(provider, current_dropdown):
    models = get_models(provider)
    return models, models[0] if models else "Nessuno"

# Aggiorna modelli per provider opzionale
def update_models_optional(provider, current_dropdown):
    if provider == "Nessuno":
        return ["Nessuno"], "Nessuno"
    models = get_models(provider)
    return models, models[0] if models else "Nessuno"

# Crea l'interfaccia Gradio
with gr.Blocks(title="TOKEN ANALYSIS APP") as demo:
    # Titolo
    gr.Markdown("# 🔍 TOKEN ANALYSIS APP\nAnalizza e confronta i costi dei token per diversi modelli LLM")
    
    with gr.Tabs():
        # Prima tab: Token Counter con Prezzi
        with gr.Tab("💲 Token Counter"):
            gr.Markdown("## 💰 Calcolo del Costo dei Token\nAnalizza i tuoi testi e confronta i costi tra diversi modelli")
            
            # Sezione input/output
            with gr.Row():
                input_text = gr.Textbox(label="📝 Input Text", lines=5, placeholder="Inserisci qui il testo di input...")
                output_text = gr.Textbox(label="📄 Output Text", lines=5, placeholder="Inserisci qui il testo di output...")
            
            # Sezione selezione modelli
            with gr.Row():
                with gr.Column(scale=1):
                    # Dropdown del primo provider/modello
                    gr.Markdown("### 🔵 Modello 1")
                    provider_dropdown = gr.Dropdown(
                        choices=unique_providers,
                        label="Provider LLM 1",
                        value=initial_provider,
                        interactive=True
                    )
                    
                    model_dropdown = gr.Dropdown(
                        choices=initial_models,
                        label="Modello LLM 1",
                        value=initial_model,
                        interactive=True
                    )
                
                with gr.Column(scale=1):
                    # Dropdown del secondo provider/modello per confronto
                    gr.Markdown("### 🟠 Modello 2 (opzionale)")
                    provider_dropdown_compare = gr.Dropdown(
                        choices=["Nessuno"] + unique_providers,
                        label="Provider LLM 2",
                        value="Nessuno",
                        interactive=True
                    )
                    
                    model_dropdown_compare = gr.Dropdown(
                        choices=["Nessuno"],
                        label="Modello LLM 2",
                        value="Nessuno",
                        interactive=True
                    )
            
            # Bottoni azione
            with gr.Row():
                compute_btn = gr.Button("🧮 Calcola Costi", variant="primary")
                reset_btn = gr.Button("🔄 Reset", variant="secondary")
            
            # Sezione risultati
            with gr.Row():
                with gr.Column(scale=1):
                    total_token = gr.Textbox(label="🔢 Token Totali", lines=1)
                    token_info = gr.Textbox(label="📊 Informazioni Token", lines=4)
                with gr.Column(scale=2):
                    price_per_token_info = gr.Markdown(label="💸 Prezzi per Token")
            
            with gr.Row():
                token_cost_table = gr.Markdown(label="📈 Confronto Costi")
            
            # Calcolo dei costi
            compute_btn.click(
                fn=count_tokens_and_price, 
                inputs=[input_text, output_text, provider_dropdown, model_dropdown, 
                        provider_dropdown_compare, model_dropdown_compare], 
                outputs=[total_token, token_info, price_per_token_info, token_cost_table]
            )
            
            # Reset
            def reset_counter():
                return "", "", "", "", "", ""
            
            reset_btn.click(
                fn=reset_counter, 
                inputs=[], 
                outputs=[input_text, output_text, total_token, token_info, price_per_token_info, token_cost_table]
            )
        
        # Seconda tab: Token Understanding
        with gr.Tab("🧠 Token Understanding"):
            gr.Markdown("## 📚 Comprensione dei Token\nGenera esempi di testo con un numero specifico di token e confronta i costi")
            
            with gr.Row():
                ntokens = gr.Slider(minimum=10, maximum=1000000, value=100, step=10, label="🔢 Numero di Token")
            
            with gr.Row():
                with gr.Column(scale=1):
                    # Dropdown del primo provider/modello
                    gr.Markdown("### 🔵 Modello 1")
                    provider_dropdown2 = gr.Dropdown(
                        choices=unique_providers,
                        label="Provider LLM 1",
                        value=initial_provider,
                        interactive=True
                    )
                    
                    model_dropdown2 = gr.Dropdown(
                        choices=initial_models,
                        label="Modello LLM 1",
                        value=initial_model,
                        interactive=True
                    )
                
                with gr.Column(scale=1):
                    # Dropdown del secondo provider/modello per confronto
                    gr.Markdown("### 🟠 Modello 2 (opzionale)")
                    provider_dropdown2_compare = gr.Dropdown(
                        choices=["Nessuno"] + unique_providers,
                        label="Provider LLM 2",
                        value="Nessuno",
                        interactive=True
                    )
                    
                    model_dropdown2_compare = gr.Dropdown(
                        choices=["Nessuno"],
                        label="Modello LLM 2",
                        value="Nessuno",
                        interactive=True
                    )
            
            with gr.Row():
                generate_btn = gr.Button("✨ Genera Esempio", variant="primary")
            
            with gr.Row():
                with gr.Column():
                    text_example = gr.Textbox(label="📝 Esempio di Testo", lines=8)
                    tokens_count = gr.Textbox(label="🔢 Token Effettivi", lines=1)
                with gr.Column():
                    token_cost_table2 = gr.Markdown(label="💰 Costi Stimati")
            
            # Generazione esempio
            generate_btn.click(
                fn=show_text_length, 
                inputs=[ntokens, provider_dropdown2, model_dropdown2, 
                       provider_dropdown2_compare, model_dropdown2_compare], 
                outputs=[text_example, tokens_count, token_cost_table2]
            )
    
    # Funzioni per cambiare dropdown - gestione valore selezionato
    def on_provider_change(provider):
        models = get_models(provider)
        return models, models[0] if models else "Nessuno"
    
    def on_provider_compare_change(provider):
        if provider == "Nessuno":
            return ["Nessuno"], "Nessuno"
        models = get_models(provider)
        return models, models[0] if models else "Nessuno"
    
    # Gestione del cambio di provider - con output a due valori
    provider_dropdown.change(
        fn=on_provider_change,
        inputs=[provider_dropdown],
        outputs=[model_dropdown, model_dropdown]  # Aggiorna sia le opzioni che il valore
    )
    
    provider_dropdown_compare.change(
        fn=on_provider_compare_change,
        inputs=[provider_dropdown_compare],
        outputs=[model_dropdown_compare, model_dropdown_compare]  # Aggiorna sia le opzioni che il valore
    )
    
    provider_dropdown2.change(
        fn=on_provider_change,
        inputs=[provider_dropdown2],
        outputs=[model_dropdown2, model_dropdown2]  # Aggiorna sia le opzioni che il valore
    )
    
    provider_dropdown2_compare.change(
        fn=on_provider_compare_change,
        inputs=[provider_dropdown2_compare],
        outputs=[model_dropdown2_compare, model_dropdown2_compare]  # Aggiorna sia le opzioni che il valore
    )
    
    # Footer
    gr.Markdown("---\n📊 Token Analysis App - Calcola e confronta i costi dei token per diversi modelli LLM")

# Avvia l'app
if __name__ == "__main__":
    demo.launch()