# app.py
from fastapi import FastAPI, HTTPException
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from typing import List, Dict
from train import MangaRecommender, MangaDataset  # Import from train.py

app = FastAPI()

try:
    # Load model and mappings
    checkpoint = torch.load('manga_recommender.pt')
    model = MangaRecommender(
        num_users=len(checkpoint['user_mapping']),
        num_items=len(checkpoint['manga_mapping'])
    )
    model.load_state_dict(checkpoint['model_state_dict'])
    user_mapping = checkpoint['user_mapping']
    manga_mapping = checkpoint['manga_mapping']
    reverse_manga_mapping = {v: k for k, v in manga_mapping.items()}
    print("Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None
    user_mapping = {}
    manga_mapping = {}
    reverse_manga_mapping = {}

@app.get("/")
async def root():
    return {"status": "running", "model_loaded": model is not None}

@app.post("/predict")
async def predict(user_id: str, top_k: int = 10):
    if model is None:
        raise HTTPException(status_code=500, detail="Model not loaded")
        
    try:
        # Get user index
        user_idx = user_mapping.get(user_id)
        if user_idx is None:
            # Handle cold start
            return {"error": "User not found"}
        
        # Get predictions
        model.eval()
        with torch.no_grad():
            user_tensor = torch.tensor([user_idx])
            predictions = model.predict(user_tensor)
            scores, indices = torch.topk(predictions[0], k=top_k)
            
            # Convert back to manga IDs
            manga_ids = [reverse_manga_mapping[idx.item()] for idx in indices]
            scores = scores.tolist()
            
            return {
                "manga_ids": manga_ids,
                "scores": scores
            }
    except Exception as e:
        raise HTTPException(
            status_code=500, 
            detail=f"Prediction error: {str(e)}"
        )

@app.post("/update")
async def update_model(ratings: List[Dict]):
    if model is None:
        raise HTTPException(status_code=500, detail="Model not loaded")
        
    try:
        # Convert ratings to training format
        df = pd.DataFrame(ratings)
        df['user_idx'] = df['user_id'].map(user_mapping)
        df['manga_idx'] = df['manga_id'].map(manga_mapping)
        
        # Create dataset
        dataset = MangaDataset(df)
        loader = DataLoader(dataset, batch_size=64)
        
        # Update model
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        criterion = nn.MSELoss()
        
        model.train()
        total_loss = 0
        for user, item, rating in loader:
            optimizer.zero_grad()
            pred = model(user, item)
            loss = criterion(pred, rating)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            
        # Save updated model
        torch.save({
            'model_state_dict': model.state_dict(),
            'user_mapping': user_mapping,
            'manga_mapping': manga_mapping
        }, 'manga_recommender.pt')
        
        return {
            "message": "Model updated successfully",
            "average_loss": total_loss / len(loader)
        }
    except Exception as e:
        raise HTTPException(
            status_code=500, 
            detail=f"Update error: {str(e)}"
        )

@app.get("/model-info")
async def model_info():
    if model is None:
        raise HTTPException(status_code=500, detail="Model not loaded")
        
    return {
        "num_users": len(user_mapping),
        "num_manga": len(manga_mapping),
        "embedding_size": model.user_factors.embedding_dim
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
    raise HTTPException(status_code=500, detail=str(e))