from fastapi import FastAPI
from fastapi.responses import JSONResponse
import firebase_admin
from firebase_admin import credentials, firestore
from transformers import pipeline
from pydantic import BaseModel
import os
from huggingface_hub import login

# Load Firebase
cred = credentials.Certificate("firebase_config.json")
firebase_admin.initialize_app(cred)
db = firestore.client()

# Đăng nhập vào Hugging Face (nếu cần)
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    login(HF_TOKEN)

# Load AI Model
ai_model = pipeline("text2text-generation", model="google/flan-t5-base")

app = FastAPI()

# Request Models
class UpdateBiosRequest(BaseModel):
    user_id: str
    bios: str

class FocusHistoryRequest(BaseModel):
    user_id: str
    time_start: str
    total_time: int

class AIRequest(BaseModel):
    user_id: str

# API cập nhật bios
@app.post("/update_bios")
async def update_bios(request: UpdateBiosRequest):
    user_ref = db.collection("user_profiles").document(request.user_id)
    user_ref.set({"bios": request.bios}, merge=True)
    return {"message": "Cập nhật bios thành công"}

# API thêm lịch sử focus
@app.post("/add_focus_history")
async def add_focus_history(request: FocusHistoryRequest):
    user_ref = db.collection("user_profiles").document(request.user_id)
    user_doc = user_ref.get()
    data = user_doc.to_dict() or {}
    focus_history = data.get("focus_history", [])
    focus_history.append({"time_start": request.time_start, "total_time": request.total_time})
    user_ref.set({"focus_history": focus_history}, merge=True)
    return {"message": "Thêm lịch sử focus thành công"}

# API lấy dữ liệu người dùng
@app.get("/get_user_data")
async def get_user_data(user_id: str):
    user_doc = db.collection("user_profiles").document(user_id).get()
    data = user_doc.to_dict() or {}
    return {
        "bios": data.get("bios", "Chưa có bios."),
        "focus_history": data.get("focus_history", [])
    }

# API AI tư vấn
@app.post("/ai_personal_advice")
async def ai_personal_advice(request: AIRequest):
    user_doc = db.collection("user_profiles").document(request.user_id).get()
    data = user_doc.to_dict() or {}
    bios = data.get("bios", "Chưa có bios.")
    focus_history = data.get("focus_history", [])
    focus_text = "\n".join([f"- {f['time_start']}: {f['total_time']} phút" for f in focus_history])
    
    prompt = f"""
    Thông tin người dùng:
    - Bios: {bios}
    - Lịch sử focus:
    {focus_text}
    
    Hãy tư vấn cách cải thiện hiệu suất làm việc dựa trên thông tin trên.
    """
    
    response = ai_model(prompt, max_length=200)
    return {"advice": response[0]["generated_text"]}

# Trang chủ
@app.get("/")
async def home():
    return JSONResponse(content={"message": "Welcome to the Recommendation API!"})