import os

import dotenv
import pandas as pd
from datasets import load_dataset
from fasthtml.common import *

dotenv.load_dotenv()

app, rt = fast_app()
from huggingface_hub import login, whoami

login(token=os.environ.get("HF_TOKEN"))

hf_user = whoami(os.environ.get("HF_TOKEN"))["name"]
HF_REPO_ID = f"{hf_user}/zotero-answer-ai-articles"


@rt("/")
async def get():
    """Dataset viewer home page"""
    return Titled(
        "Zotero Dataset Viewer",
        Form(
            Input(id="dataset_subset", placeholder="Enter dataset subset"),
            Button("Load Dataset"),
            hx_post="/load_dataset",
            hx_target="#dataset_content",
        ),
        Div(id="dataset_content"),
    )


@rt("/load_dataset")
async def post(dataset_subset: str):
    """Load and display dataset"""
    try:
        dataset = load_dataset(HF_REPO_ID, dataset_subset, split="train")
        df = pd.DataFrame(dataset[:10])  # Load first 10 rows
        table = df.to_html(classes="table", index=False)

        return Div(
            H2(f"Dataset: {HF_REPO_ID}, Subset: {dataset_subset}"),
            P(f"Number of rows: {len(dataset)}"),
            P(f"Columns: {', '.join(df.columns)}"),
            Div(NotStr(table), cls="table-responsive"),
        )

    except Exception as e:
        return Div(f"Error loading dataset: {str(e)}")


serve()