import os
import json
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from datasets import load_dataset
import uvicorn

app = FastAPI(title="Human Essence Emotional Flow Visualizer")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = "wop/Human-Essence-Dataset"

def fetch_dataset(force=False):
    download_mode = "force_redownload" if force else None
    ds = load_dataset(DATASET_REPO, token=HF_TOKEN, split="train", download_mode=download_mode)
    return ds.to_list()


print("Loading dataset from Hugging Face...")
try:
    dataset = fetch_dataset()
    print(f"Loaded {len(dataset)} entries successfully!")
except Exception as e:
    print(f"Error loading dataset: {e}")
    dataset = []


@app.post("/api/refresh")
def refresh_data():
    global dataset
    try:
        dataset = fetch_dataset(force=True)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Refresh failed: {e}")
    print(f"Refreshed dataset: {len(dataset)} entries")
    return {"ok": True, "total_entries": len(dataset)}


@app.get("/api/dataset")
def get_dataset(limit: int = 100, offset: int = 0):
    """Get dataset entries with pagination"""
    total = len(dataset)
    end_idx = min(offset + limit, total)
    entries = dataset[offset:end_idx]
    
    return {
        "entries": entries,
        "total": total,
        "offset": offset,
        "limit": limit
    }


@app.get("/api/entry/{index}")
def get_entry(index: int):
    """Get a specific entry by index"""
    if index < 0 or index >= len(dataset):
        raise HTTPException(status_code=404, detail="Entry not found")
    
    return dataset[index]


@app.get("/api/stats")
def get_stats():
    """Get dataset statistics"""
    total_entries = len(dataset)
    entries_with_text = sum(1 for e in dataset if e.get("text"))
    entries_with_emotions = sum(1 for e in dataset if e.get("emotional_flow"))
    
    # Count unique emotions
    emotion_counts = {}
    for entry in dataset:
        for flow in entry.get("emotional_flow", []):
            label = flow.get("label", "unknown")
            emotion_counts[label] = emotion_counts.get(label, 0) + 1
    
    return {
        "total_entries": total_entries,
        "entries_with_text": entries_with_text,
        "entries_with_emotions": entries_with_emotions,
        "emotion_distribution": emotion_counts
    }


class NoCacheStaticFiles(StaticFiles):
    async def get_response(self, path, scope):
        response = await super().get_response(path, scope)
        response.headers["Cache-Control"] = "no-cache, must-revalidate"
        return response


app.mount("/", NoCacheStaticFiles(directory="static", html=True), name="static")

if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)