import gradio as gr
import random
import json
import fastapi
from fastapi import FastAPI
import os
import argilla as rg

# Initialize Argilla client
client = rg.Argilla(
    api_url=os.getenv("ARGILLA_API_URL", ""),
    api_key=os.getenv("ARGILLA_API_KEY", "")
)

# List of countries to check
COUNTRIES = [
    "MEX", "ARG", "COL", "CHL", "PER", "ESP", "BRA", 
    "VEN", "ECU", "BOL", "PRY", "URY", "CRI", "PAN", 
    "DOM", "GTM", "HND", "SLV", "NIC", "CUB"
]

def count_answers_per_space(country: str):
    """
    Count the number of answers for a specific country's answering space.
    
    Args:
        country: Country code (e.g., "CHL")
    
    Returns:
        Dictionary with statistics about answers in the space
    """
    # Convert country code to full country name
    country_mapping = {
        "MEX": "Mexico", "ARG": "Argentina", "COL": "Colombia", 
        "CHL": "Chile", "PER": "Peru", "ESP": "Spain", 
        "BRA": "Brazil", "VEN": "Venezuela", "ECU": "Ecuador", 
        "BOL": "Bolivia", "PRY": "Paraguay", "URY": "Uruguay", 
        "CRI": "Costa Rica", "PAN": "Panama", "DOM": "Dominican Republic", 
        "GTM": "Guatemala", "HND": "Honduras", "SLV": "El Salvador", 
        "NIC": "Nicaragua", "CUB": "Cuba"
    }
    
    full_country_name = country_mapping.get(country, country)
    dataset_name = f"{full_country_name}_responder_preguntas"
    
    try:
        dataset = client.datasets(dataset_name)
        
        # Get all records with their responses
        records = dataset.records(with_responses=True)
        
        # Count total questions
        total_questions = len(list(records))
        
        # Count answered questions
        answered_questions = 0
        total_answers = 0
        answers_per_user = {}
        
        for record in records:
            responses = record.responses.get("text", [])
            if responses:
                answered_questions += 1
                total_answers += len(responses)
                
                # Count per user
                for response in responses:
                    user_id = str(response.user_id)
                    if user_id in answers_per_user:
                        answers_per_user[user_id] += 1
                    else:
                        answers_per_user[user_id] = 1
        
        percentage_complete = (answered_questions / total_questions * 100) if total_questions > 0 else 0
        
        return {
            "name": full_country_name,
            "total_questions": total_questions,
            "answered_questions": answered_questions,
            "total_answers": total_answers,
            "percent": round(percentage_complete, 2),
            "documents": total_answers * 10  # Approximation of document count
        }
        
    except Exception as e:
        # If space doesn't exist, return zero values
        print(f"No dataset found for {dataset_name}: {e}")
        return {
            "name": full_country_name,
            "total_questions": 0,
            "answered_questions": 0,
            "total_answers": 0,
            "percent": 0,
            "documents": 0
        }

# Create a FastAPI app
app = FastAPI()

# Route to serve the map visualization
@app.get("/d3-map")
async def serve_map():
    # Generate data for each country by querying Argilla
    country_data = {}
    for country_code in COUNTRIES:
        country_data[country_code] = count_answers_per_space(country_code)
    
    # Convert to JSON for JavaScript
    country_data_json = json.dumps(country_data)
    
    # Replace the placeholder with actual data
    with open('template.txt', 'r') as f:
        html_template = f.read()
    
    html_content = html_template.replace("COUNTRY_DATA_PLACEHOLDER", country_data_json)
    
    return fastapi.responses.HTMLResponse(content=html_content)

# Create a simple Gradio interface with an iframe
def create_iframe():
    # Add a random parameter to force reload
    random_param = random.randint(1, 10000)
    return '<iframe src="/d3-map?t={}" style="width:100%; height:650px; border:none;"></iframe>'.format(random_param)

# Create the Gradio blocks
with gr.Blocks(theme=gr.themes.Soft(primary_hue="pink", secondary_hue="purple")) as demo:
    gr.Markdown("# Mapa anotación")
    
    iframe_output = gr.HTML(create_iframe())
    
    # Refresh button to generate new random data
    def refresh():
        return create_iframe()
    
    gr.Button("Actualizar Datos").click(fn=refresh, outputs=iframe_output)

# Mount the Gradio app to the FastAPI app
gr.mount_gradio_app(app, demo, path="/")

# Start the server
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)