from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import login
import os
from dotenv import load_dotenv

from utils import (
    get_data,
    format_data,
    save_data,
    make_prediction
)

# Initialize FastAPI
app = FastAPI()

# Load environment variables
load_dotenv()

class ScrapeRequest(BaseModel):
    start_year: int
    end_year: int

@app.post("/scrape/")
async def scrape(request: ScrapeRequest):
    try:
        # Login to Hugging Face
        login(token=os.getenv("HF_TOKEN"))
        
        results = []
        for year in range(request.start_year, request.end_year + 1):
            try:
                stats = get_data(year)
                formatted_stats = format_data(stats, year)
                await save_data(formatted_stats, year, predict_callback=make_prediction)
                results.append({
                    "year": year,
                    "status": "success",
                    "message": f"Data for year {year-1}-{year} saved successfully"
                })
            except Exception as e:
                results.append({
                    "year": year,
                    "status": "error",
                    "message": str(e)
                })
        
        return {"results": results}
    
    except Exception as e:
        return {"error": str(e)}

@app.post("/predict/")
async def predict():
    return await make_prediction()

@app.get("/")
async def root():
    return {"message": "Welcome to the Curry Model API!"}