Update main.py

main.py

@@ -1,12 +1,14 @@
 import os
 import secrets
 import requests
-from fastapi import FastAPI, Depends, HTTPException, Header
+from fastapi import FastAPI, Depends, HTTPException, Header, Request
 from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, func
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, Session
 from pydantic import BaseModel
 from dotenv import load_dotenv
+from fastapi.templating import Jinja2Templates
+from fastapi.staticfiles import StaticFiles
 
 # Load environment variables from .env file
 load_dotenv()
@@ -27,7 +29,7 @@ engine = create_engine(DATABASE_URL)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 Base = declarative_base()
 
-# User model
+# Updated User model with credits field
 class User(Base):
     __tablename__ = "users"
     id = Column(Integer, primary_key=True, index=True)
@@ -35,12 +37,17 @@ class User(Base):
     hashed_password = Column(String(128), nullable=False)
     api_key = Column(String(64), unique=True, index=True, nullable=False)
     is_admin = Column(Boolean, default=False)
+    credits = Column(Integer, default=0)
     created_at = Column(DateTime(timezone=True), server_default=func.now())
 
 # Create tables
 Base.metadata.create_all(bind=engine)
 
-app = FastAPI(title="API Key Generator & Proxy Service")
+app = FastAPI(title="API Key Platform")
+
+# Mount static files and templates
+app.mount("/static", StaticFiles(directory="static"), name="static")
+templates = Jinja2Templates(directory="templates")
 
 # Dependency: Database session
 def get_db():
@@ -64,6 +71,7 @@ class UserOut(BaseModel):
     username: str
     api_key: str
     is_admin: bool
+    credits: int
 
     class Config:
         orm_mode = True
@@ -103,7 +111,7 @@ def list_users(current_user: User = Depends(get_current_user), db: Session = Dep
     users = db.query(User).all()
     return users
 
-# 4. Endpoint to access the main API (proxying request using main API key)
+# 4. Proxy Endpoint to access the main API
 class RequestPayload(BaseModel):
     prompt: str
 
@@ -121,3 +129,37 @@ def generate_image(payload: RequestPayload, current_user: User = Depends(get_cur
     if response.status_code != 200:
         raise HTTPException(status_code=response.status_code, detail="Error from main API")
     return response.json()
+
+# 5. New endpoint for users to test their API key
+@app.get("/user/test_api")
+def test_api(current_user: User = Depends(get_current_user)):
+    return {"message": "API is working", "username": current_user.username, "credits": current_user.credits}
+
+# 6. New endpoint for admin to add credits to a user account
+class CreditPayload(BaseModel):
+    username: str
+    credits: int
+
+@app.post("/admin/add_credit")
+def add_credit(payload: CreditPayload, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
+    if not current_user.is_admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
+    user = db.query(User).filter(User.username == payload.username).first()
+    if not user:
+        raise HTTPException(status_code=404, detail="User not found")
+    user.credits += payload.credits
+    db.commit()
+    db.refresh(user)
+    return {"message": f"Added {payload.credits} credits to user {user.username}. Total credits: {user.credits}"}
+
+# 7. Render Admin Panel UI
+@app.get("/admin/ui")
+def admin_ui(request: Request, current_user: User = Depends(get_current_user)):
+    if not current_user.is_admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
+    return templates.TemplateResponse("admin.html", {"request": request})
+
+# 8. Render User Panel UI
+@app.get("/user/ui")
+def user_ui(request: Request, current_user: User = Depends(get_current_user)):
+    return templates.TemplateResponse("user.html", {"request": request})