backend/main.py

import os
import time
import random
import smtplib
import bcrypt
from datetime import datetime, timedelta
from email.mime.text import MIMEText
from typing import List, Optional

from fastapi import FastAPI, Depends, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from fastapi.staticfiles import StaticFiles
from jose import jwt, JWTError
from pydantic import BaseModel, EmailStr, Field
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime
from sqlalchemy.orm import sessionmaker, Session, DeclarativeBase
from dotenv import load_dotenv

load_dotenv()

# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════════════════════

SECRET_KEY = os.getenv("SECRET_KEY", "super_secret_dev_key_change_me_123456789")
ALGORITHM = os.getenv("ALGORITHM", "HS256")
ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "30"))
REFRESH_TOKEN_EXPIRE_DAYS = int(os.getenv("REFRESH_TOKEN_EXPIRE_DAYS", "7"))

DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./users.db")

SMTP_SERVER = os.getenv("SMTP_SERVER", "")
SMTP_PORT = int(os.getenv("SMTP_PORT", "587"))
SMTP_USER = os.getenv("SMTP_USER", "")
SMTP_PASSWORD = os.getenv("SMTP_PASSWORD", "")

# ══════════════════════════════════════════════════════════════════════════════
# DATABASE
# ══════════════════════════════════════════════════════════════════════════════

connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}

engine = create_engine(DATABASE_URL, connect_args=connect_args)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)


class Base(DeclarativeBase):
    pass


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# ══════════════════════════════════════════════════════════════════════════════
# MODELS
# ══════════════════════════════════════════════════════════════════════════════

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, index=True, nullable=False)
    email = Column(String, unique=True, index=True, nullable=False)
    password = Column(String, nullable=False)
    is_active = Column(Boolean, default=True)
    role = Column(String, default="user")
    created_at = Column(DateTime, default=datetime.utcnow)
    is_verified = Column(Boolean, default=False)
    verification_code = Column(String, nullable=True)
    reset_code = Column(String, nullable=True)


class RevokedToken(Base):
    __tablename__ = "revoked_tokens"

    id = Column(Integer, primary_key=True, index=True)
    token = Column(String, unique=True, index=True, nullable=False)
    revoked_at = Column(DateTime, default=datetime.utcnow)
    expires_at = Column(DateTime, nullable=False)


# ══════════════════════════════════════════════════════════════════════════════
# SCHEMAS
# ══════════════════════════════════════════════════════════════════════════════

class RegisterInput(BaseModel):
    username: str = Field(..., min_length=3, max_length=50)
    email: EmailStr
    password: str = Field(..., min_length=8)


class TokenOutput(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str = "bearer"


class RefreshInput(BaseModel):
    refresh_token: str


class UserOutput(BaseModel):
    id: int
    username: str
    email: str
    role: str
    is_active: bool
    is_verified: bool

    class Config:
        from_attributes = True


class ProfileUpdateInput(BaseModel):
    username: Optional[str] = Field(None, min_length=3, max_length=50)
    email: Optional[EmailStr] = None
    password: Optional[str] = Field(None, min_length=8)


class VerifyEmailInput(BaseModel):
    email: EmailStr
    code: str = Field(..., min_length=6, max_length=6, description="6-digit verification code")


class ForgotPasswordRequest(BaseModel):
    email: EmailStr


class PasswordResetInput(BaseModel):
    email: EmailStr
    code: str = Field(..., min_length=6, max_length=6, description="6-digit reset code")
    new_password: str = Field(..., min_length=8, description="Must be at least 8 characters")


# ══════════════════════════════════════════════════════════════════════════════
# PASSWORD HASHING
# ══════════════════════════════════════════════════════════════════════════════

def hash_password(password: str) -> str:
    salt = bcrypt.gensalt()
    return bcrypt.hashpw(password.encode("utf-8"), salt).decode("utf-8")


def check_password(password: str, hashed_password: str) -> bool:
    try:
        return bcrypt.checkpw(password.encode("utf-8"), hashed_password.encode("utf-8"))
    except Exception:
        return False


# ══════════════════════════════════════════════════════════════════════════════
# JWT TOKENS
# ══════════════════════════════════════════════════════════════════════════════

def make_access_token(user_id: int, email: str) -> str:
    expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    data = {"sub": str(user_id), "email": email, "type": "access", "exp": expire}
    return jwt.encode(data, SECRET_KEY, algorithm=ALGORITHM)


def make_refresh_token(user_id: int) -> str:
    expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
    data = {"sub": str(user_id), "type": "refresh", "exp": expire}
    return jwt.encode(data, SECRET_KEY, algorithm=ALGORITHM)


# ══════════════════════════════════════════════════════════════════════════════
# OTP GENERATOR
# ══════════════════════════════════════════════════════════════════════════════

def generate_otp() -> str:
    return "".join(random.choices("0123456789", k=6))


# ══════════════════════════════════════════════════════════════════════════════
# EMAIL SENDING (SMTP / Console Fallback)
# ══════════════════════════════════════════════════════════════════════════════

def send_email(to_email: str, subject: str, body: str) -> bool:
    if not all([SMTP_SERVER, SMTP_USER, SMTP_PASSWORD]):
        print(f"\n┌────────────────────────────────────────────────────────┐")
        print(f"│ [SIMULATION] Email Sent To: {to_email:<27} │")
        print(f"│ Subject: {subject:<44} │")
        print(f"│ Code/OTP: {body:<44} │")
        print(f"└────────────────────────────────────────────────────────┘\n")
        return True

    try:
        msg = MIMEText(body)
        msg["Subject"] = subject
        msg["From"] = SMTP_USER
        msg["To"] = to_email

        with smtplib.SMTP(SMTP_SERVER, SMTP_PORT, timeout=5) as server:
            server.starttls()
            server.login(SMTP_USER, SMTP_PASSWORD)
            server.send_message(msg)
        return True
    except Exception as e:
        print(f"\n[SMTP ERROR] Sending failed: {e}")
        print(f"┌────────────────────────────────────────────────────────┐")
        print(f"│ [SIMULATION FALLBACK] Email To: {to_email:<22} │")
        print(f"│ Subject: {subject:<44} │")
        print(f"│ Code/OTP: {body:<44} │")
        print(f"└────────────────────────────────────────────────────────┘\n")
        return False


# ══════════════════════════════════════════════════════════════════════════════
# CURRENT USER DEPENDENCY
# ══════════════════════════════════════════════════════════════════════════════

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/login")


def get_current_user(
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
) -> User:
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )

    revoked = db.query(RevokedToken).filter(RevokedToken.token == token).first()
    if revoked:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Session has expired or you logged out. Please sign in again.",
            headers={"WWW-Authenticate": "Bearer"},
        )

    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        if payload.get("type") != "access":
            raise HTTPException(status_code=401, detail="Invalid token type")
        user_id: str = payload.get("sub")
        if user_id is None:
            raise credentials_exception
    except JWTError:
        raise credentials_exception

    user = db.query(User).filter(User.id == int(user_id)).first()
    if user is None:
        raise HTTPException(status_code=404, detail="User not found")
    if not user.is_active:
        raise HTTPException(status_code=400, detail="User account is inactive")
    if not user.is_verified:
        raise HTTPException(status_code=400, detail="User email is not verified yet")

    return user


# ══════════════════════════════════════════════════════════════════════════════
# PASSWORD VALIDATION HELPER
# ══════════════════════════════════════════════════════════════════════════════

def validate_password_strength(password: str):
    if len(password) < 8:
        raise HTTPException(status_code=400, detail="Password must be at least 8 characters long.")
    if not any(char.isdigit() for char in password):
        raise HTTPException(status_code=400, detail="Password must contain at least one number.")
    if not any(char.isalpha() for char in password):
        raise HTTPException(status_code=400, detail="Password must contain at least one letter.")


# ══════════════════════════════════════════════════════════════════════════════
# FASTAPI APP
# ══════════════════════════════════════════════════════════════════════════════

Base.metadata.create_all(bind=engine)

app = FastAPI(
    title="SMTP-Enabled Authentication System",
    description="A beginner-friendly secure authentication system with modular routes and email verification.",
    version="1.0.0",
)

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

# ─── RATE LIMITING & REQUEST LOGGING ──────────────────────────────────────

RATE_LIMIT_WINDOW = 60
RATE_LIMIT_MAX = 60
client_requests = {}


@app.middleware("http")
async def rate_limiting_and_logging_middleware(request: Request, call_next):
    client_ip = request.client.host if request.client else "unknown"
    now = time.time()

    if client_ip not in client_requests:
        client_requests[client_ip] = []

    client_requests[client_ip] = [t for t in client_requests[client_ip] if now - t < RATE_LIMIT_WINDOW]

    if len(client_requests[client_ip]) >= RATE_LIMIT_MAX:
        return JSONResponse(
            status_code=429,
            content={"detail": "Too many requests. Please slow down and try again in a minute."},
        )

    client_requests[client_ip].append(now)
    start_time = time.time()

    try:
        response = await call_next(request)
    except Exception as exc:
        log_line = f"{datetime.utcnow().isoformat()} - {client_ip} - {request.method} {request.url.path} - ERROR: {str(exc)}\n"
        with open("requests.log", "a") as log_file:
            log_file.write(log_line)
        return JSONResponse(
            status_code=500,
            content={"detail": "An internal server error occurred. Please contact the administrator."},
        )

    duration = time.time() - start_time
    log_line = (
        f"{datetime.utcnow().isoformat()} - IP: {client_ip} - "
        f"{request.method} {request.url.path} - Status: {response.status_code} - "
        f"Took: {duration:.4f}s\n"
    )
    with open("requests.log", "a") as log_file:
        log_file.write(log_line)

    return response


# ══════════════════════════════════════════════════════════════════════════════
# AUTH ROUTES
# ══════════════════════════════════════════════════════════════════════════════

@app.post("/auth/register", response_model=UserOutput, status_code=201)
def register(data: RegisterInput, db: Session = Depends(get_db)):
    validate_password_strength(data.password)

    if db.query(User).filter(User.username == data.username).first():
        raise HTTPException(status_code=400, detail="Username is already taken")
    if db.query(User).filter(User.email == data.email).first():
        raise HTTPException(status_code=400, detail="Email is already registered")

    otp_code = generate_otp()

    new_user = User(
        username=data.username,
        email=data.email,
        password=hash_password(data.password),
        is_verified=False,
        verification_code=otp_code,
    )

    db.add(new_user)
    db.commit()
    db.refresh(new_user)

    send_email(to_email=new_user.email, subject="Email Verification Code", body=otp_code)

    return new_user


@app.post("/auth/verify-email")
def verify_email(data: VerifyEmailInput, db: Session = Depends(get_db)):
    user = db.query(User).filter(User.email == data.email).first()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    if user.is_verified:
        return {"detail": "Email already verified. You can sign in."}

    if user.verification_code != data.code:
        raise HTTPException(status_code=400, detail="Invalid verification code")

    user.is_verified = True
    user.verification_code = None
    db.commit()

    return {"detail": "Email verified successfully! You can now sign in."}


@app.post("/auth/login", response_model=TokenOutput)
def login(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
    user = db.query(User).filter(
        (User.username == form_data.username) | (User.email == form_data.username)
    ).first()

    if not user or not check_password(form_data.password, user.password):
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect username or password")

    if not user.is_active:
        raise HTTPException(status_code=400, detail="User account is deactivated")

    if not user.is_verified:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Email is not verified. Please verify your email first.",
        )

    access = make_access_token(user.id, user.email)
    refresh = make_refresh_token(user.id)

    return TokenOutput(access_token=access, refresh_token=refresh)


@app.post("/auth/forgot-password")
def forgot_password(data: ForgotPasswordRequest, db: Session = Depends(get_db)):
    user = db.query(User).filter(User.email == data.email).first()
    if not user:
        raise HTTPException(status_code=404, detail="Email not found")

    otp_code = generate_otp()
    user.reset_code = otp_code
    db.commit()

    send_email(to_email=user.email, subject="Password Reset Request Code", body=otp_code)

    return {"detail": "Verification code has been sent to your email."}


@app.post("/auth/reset-password")
def reset_password(data: PasswordResetInput, db: Session = Depends(get_db)):
    validate_password_strength(data.new_password)

    user = db.query(User).filter(User.email == data.email).first()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    if user.reset_code != data.code:
        raise HTTPException(status_code=400, detail="Invalid password reset code")

    user.password = hash_password(data.new_password)
    user.reset_code = None
    db.commit()

    return {"detail": "Password has been reset successfully."}


@app.post("/auth/refresh", response_model=TokenOutput)
def refresh(data: RefreshInput, db: Session = Depends(get_db)):
    try:
        payload = jwt.decode(data.refresh_token, SECRET_KEY, algorithms=[ALGORITHM])
        if payload.get("type") != "refresh":
            raise HTTPException(status_code=401, detail="Invalid token type")
        user_id = int(payload.get("sub"))
        exp_timestamp = payload.get("exp")
        expires_at = datetime.utcfromtimestamp(exp_timestamp)
    except JWTError:
        raise HTTPException(status_code=401, detail="Refresh token is expired or invalid")

    revoked = db.query(RevokedToken).filter(RevokedToken.token == data.refresh_token).first()
    if revoked:
        raise HTTPException(status_code=401, detail="Refresh token has been revoked")

    user = db.query(User).filter(User.id == user_id).first()
    if not user or not user.is_active or not user.is_verified:
        raise HTTPException(status_code=401, detail="User not found, inactive, or unverified")

    db_revoked = RevokedToken(token=data.refresh_token, expires_at=expires_at)
    db.add(db_revoked)

    access = make_access_token(user.id, user.email)
    new_refresh = make_refresh_token(user.id)

    db.commit()
    return TokenOutput(access_token=access, refresh_token=new_refresh)


@app.post("/auth/logout")
def logout(
    data: RefreshInput = None,
    access_token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    try:
        payload = jwt.decode(access_token, SECRET_KEY, algorithms=[ALGORITHM])
        exp = datetime.utcfromtimestamp(payload.get("exp"))
    except JWTError:
        exp = datetime.utcnow() + timedelta(minutes=30)

    if not db.query(RevokedToken).filter(RevokedToken.token == access_token).first():
        db.add(RevokedToken(token=access_token, expires_at=exp))

    if data and data.refresh_token:
        try:
            r_payload = jwt.decode(data.refresh_token, SECRET_KEY, algorithms=[ALGORITHM])
            r_exp = datetime.utcfromtimestamp(r_payload.get("exp"))
            if not db.query(RevokedToken).filter(RevokedToken.token == data.refresh_token).first():
                db.add(RevokedToken(token=data.refresh_token, expires_at=r_exp))
        except JWTError:
            pass

    db.commit()
    return {"detail": "Successfully logged out and session revoked"}


# ══════════════════════════════════════════════════════════════════════════════
# USER ROUTES
# ══════════════════════════════════════════════════════════════════════════════

@app.get("/me", response_model=UserOutput)
def get_profile(current_user: User = Depends(get_current_user)):
    return current_user


@app.put("/me", response_model=UserOutput)
def update_profile(
    data: ProfileUpdateInput,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    if data.username and data.username != current_user.username:
        existing = db.query(User).filter(User.username == data.username).first()
        if existing:
            raise HTTPException(status_code=400, detail="Username is already taken")
        current_user.username = data.username

    if data.email and data.email != current_user.email:
        existing = db.query(User).filter(User.email == data.email).first()
        if existing:
            raise HTTPException(status_code=400, detail="Email is already registered")
        current_user.email = data.email

    if data.password:
        validate_password_strength(data.password)
        current_user.password = hash_password(data.password)

    db.commit()
    db.refresh(current_user)
    return current_user


@app.delete("/me", status_code=200)
def delete_account(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    db.delete(current_user)
    db.commit()
    return {"detail": "Account deleted successfully"}


@app.get("/admin/users", response_model=List[UserOutput])
def get_all_users_admin(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    if current_user.role != "admin":
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Forbidden: Admin access required",
        )
    users = db.query(User).all()
    return users


# ══════════════════════════════════════════════════════════════════════════════
# HEALTH & HOME
# ══════════════════════════════════════════════════════════════════════════════

@app.get("/api")
def home():
    return {
        "status": "online",
        "message": "Welcome to the SMTP Auth API! Go to /docs to test all endpoints.",
    }


@app.get("/health")
def health():
    return {"status": "ok"}


# ─── SERVE FRONTEND STATIC FILES ──────────────────────────────────────────
if os.path.exists("static"):
    app.mount("/", StaticFiles(directory="static", html=True), name="static")


@app.exception_handler(404)
async def spa_fallback_404_handler(request: Request, exc: Exception):
    path = request.url.path
    if not path.startswith("/auth") and not path.startswith("/admin") and path != "/api":
        if os.path.exists("static/index.html"):
            with open("static/index.html", "r") as f:
                return HTMLResponse(content=f.read(), status_code=200)
    return JSONResponse(status_code=404, content={"detail": "Not Found"})