app.py

from fastapi import FastAPI, HTTPException, Body, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, EmailStr
import httpx
import bcrypt
import boto3
import os
import base64
import io

app = FastAPI()

# Enable CORS so the frontend can call this API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- ĐIỀN THÔNG TIN CLOUDFLARE BẰNG HUGGING FACE SECRETS HOẶC ĐIỀN THẲNG VÀO ĐÂY ---
CF_ACCOUNT_ID = os.getenv("CF_ACCOUNT_ID", "ID_CỦA_SẾP")
CF_D1_ID = os.getenv("CF_D1_ID", "ID_DATABASE_D1_VỪA_TẠO")
CF_API_TOKEN = os.getenv("CF_API_TOKEN", "TOKEN_BẢO_MẬT_CỦA_SẾP") 

CF_API_BASE = f"https://api.cloudflare.com/client/v4/accounts/{CF_ACCOUNT_ID}/d1/database/{CF_D1_ID}/query"

# Lấy chìa khóa R2 từ két sắt
R2_ACCESS_KEY = os.getenv("R2_ACCESS_KEY")
R2_SECRET_KEY = os.getenv("R2_SECRET_KEY")
R2_ENDPOINT = os.getenv("R2_ENDPOINT")
BUCKET_NAME = "text2-bucket" # Tên cái bucket R2 sếp tạo trên Cloudflare
R2_PUBLIC_URL = "https://data.text2.co" # Link gốc để thiên hạ vô xem ảnh sếp

# Cấu hình xe chở hàng boto3
s3_client = boto3.client(
    's3',
    endpoint_url=R2_ENDPOINT,
    aws_access_key_id=R2_ACCESS_KEY,
    aws_secret_access_key=R2_SECRET_KEY,
    region_name='auto' # R2 xài region 'auto'
) if R2_ACCESS_KEY else None

# Hàm vạn năng up file lên kho khủng R2
async def upload_to_r2(file: UploadFile, folder: str, filename: str):
    try:
        # Đường dẫn trong kho: ví dụ avatars/hoanhatanh.png
        object_name = f"{folder}/{filename}" 
        
        if not s3_client:
            print("Chưa cấu hình R2")
            return None

        s3_client.upload_fileobj(
            file.file, 
            BUCKET_NAME, 
            object_name,
            ExtraArgs={"ContentType": file.content_type}
        )
        
        # Up xong thì trả về link này để lưu vô cột avatar TEXT trong D1
        link_avatar_xin = f"{R2_PUBLIC_URL}/{object_name}"
        return link_avatar_xin 
        
    except Exception as e:
        print("Lỗi up file R2:", e)
        return None

def upload_base64_to_r2(b64_str: str, folder: str, filename: str):
    try:
        if not s3_client:
            return b64_str

        if not b64_str.startswith("data:image"):
            return b64_str
            
        header, encoded = b64_str.split(",", 1)
        image_data = base64.b64decode(encoded)
        file_obj = io.BytesIO(image_data)
        
        content_type = "image/png"
        ext = "png"
        if "jpeg" in header or "jpg" in header:
            content_type = "image/jpeg"
            ext = "jpg"
            
        object_name = f"{folder}/{filename}.{ext}"
        
        s3_client.upload_fileobj(
            file_obj,
            BUCKET_NAME,
            object_name,
            ExtraArgs={"ContentType": content_type}
        )
        return f"{R2_PUBLIC_URL}/{object_name}"
    except Exception as e:
        print("Lỗi up base64 lên R2:", e)
        return b64_str

class UserRegister(BaseModel):
    email: EmailStr
    password: str
    avatar: str = None

class UserLogin(BaseModel):
    email: EmailStr
    password: str

class UserUpdate(BaseModel):
    token: str
    username: str = None
    password: str = None
    avatar: str = None

import sqlite3

async def execute_d1_query(sql: str, params: list = []):
    """
    Helper function to interact with Cloudflare D1 HTTP API.
    """
    if CF_ACCOUNT_ID == "ID_CỦA_SẾP":
        print(f"[MOCK D1 SQLite] SQL: {sql} | params: {params}")
        with sqlite3.connect("local_mock.db") as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            try:
                cursor.execute(sql, params)
                if sql.strip().upper().startswith("SELECT") or sql.strip().upper().startswith("PRAGMA"):
                    return [dict(row) for row in cursor.fetchall()]
                else:
                    conn.commit()
                    return [{"success": True}]
            except sqlite3.IntegrityError:
                raise HTTPException(status_code=400, detail="Email already registered (Database constraint)")
            except Exception as e:
                print(f"SQLite error: {e}")
                raise HTTPException(status_code=500, detail="Local database error")

    headers = {
        "Authorization": f"Bearer {CF_API_TOKEN}",
        "Content-Type": "application/json"
    }
    payload = {
        "sql": sql,
        "params": params
    }
    
    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(CF_API_BASE, headers=headers, json=payload, timeout=10.0)
        except Exception as e:
            print(f"Connection error: {e}")
            raise HTTPException(status_code=500, detail="Error connecting to database")
            
    if response.status_code != 200:
        print(f"Cloudflare API Error: {response.text}")
        raise HTTPException(status_code=500, detail="Database query error")
        
    data = response.json()
    if not data.get("success"):
        print(f"D1 Query Failed: {data}")
        raise HTTPException(status_code=500, detail="Database query failed")
        
    result = data.get("result", [{}])[0].get("results", [])
    return result

@app.on_event("startup")
async def startup_event():
    # Create the users table on startup if it doesn't exist
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        email TEXT UNIQUE NOT NULL,
        password_hash TEXT NOT NULL,
        username TEXT,
        avatar TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    );
    """
    try:
        await execute_d1_query(create_table_sql)
        # Attempt to add avatar and username columns if they are missing from previous creation
        try:
            await execute_d1_query("ALTER TABLE users ADD COLUMN avatar TEXT;")
        except Exception:
            pass # Column already exists
            
        try:
            await execute_d1_query("ALTER TABLE users ADD COLUMN username TEXT;")
        except Exception:
            pass # Column already exists
            
        print("Database initialized.")
    except Exception as e:
        print(f"Startup table creation failed (Ignore if mock): {e}")

@app.get("/")
def read_root():
    return {"status": "ok", "message": "Text2 Database API is running on Hugging Face Spaces"}

@app.post("/api/register")
async def register(user: UserRegister):
    user.email = user.email.lower().strip()
    
    # 1. Check if user already exists
    existing = await execute_d1_query("SELECT * FROM users WHERE email = ?", [user.email])
    if existing:
        raise HTTPException(status_code=400, detail="Email already registered")
        
    # 2. Hash password
    salt = bcrypt.gensalt()
    hashed_password = bcrypt.hashpw(user.password.encode('utf-8'), salt).decode('utf-8')
    
    # 3. Xử lý ảnh đại diện: Nén base64 -> Up lên R2 (nếu có cấu hình)
    avatar_url = user.avatar
    if avatar_url and avatar_url.startswith("data:image"):
        username = user.email.split('@')[0]
        avatar_url = upload_base64_to_r2(user.avatar, "avatars", username)

    # 4. Insert user into D1
    username = user.email.split('@')[0]
    await execute_d1_query(
        "INSERT INTO users (email, password_hash, username, avatar) VALUES (?, ?, ?, ?)",
        [user.email, hashed_password, username, avatar_url]
    )
    
    return {"success": True, "message": "User registered successfully"}

@app.post("/api/login")
async def login(user: UserLogin):
    user.email = user.email.lower().strip()
    
    # 1. Find user in D1
    users = await execute_d1_query("SELECT * FROM users WHERE email = ?", [user.email])

    if not users:
        raise HTTPException(status_code=401, detail="Invalid email or password")
        
    db_user = users[0]
    
    # 2. Verify password hash
    if not bcrypt.checkpw(user.password.encode('utf-8'), db_user["password_hash"].encode('utf-8')):
        raise HTTPException(status_code=401, detail="Invalid email or password")
        
    # 3. Return success and basic token (in production use real JWT)
    return {
        "success": True,
        "message": "Login successful", 
        "user": {
            "id": db_user["id"], 
            "email": db_user["email"],
            "username": db_user.get("username") or db_user["email"].split('@')[0],
            "avatar": db_user.get("avatar")
        },
        "token": f"token-{db_user['id']}" # Placeholder token
    }

@app.post("/api/update_profile")
async def update_profile(data: UserUpdate):
    if not data.token or not data.token.startswith("token-"):
        raise HTTPException(status_code=401, detail="Unauthorized")
        
    user_id = data.token.split("-")[1]
    
    users = await execute_d1_query("SELECT * FROM users WHERE id = ?", [user_id])
    if not users:
        raise HTTPException(status_code=404, detail="User not found")
        
    db_user = users[0]
    
    updates = []
    params = []
    
    if data.username:
        updates.append("username = ?")
        params.append(data.username)
        
    if data.password:
        salt = bcrypt.gensalt()
        hashed_password = bcrypt.hashpw(data.password.encode('utf-8'), salt).decode('utf-8')
        updates.append("password_hash = ?")
        params.append(hashed_password)
        
    if data.avatar and data.avatar.startswith("data:image"):
        username = data.username or db_user.get("username") or db_user["email"].split('@')[0]
        avatar_url = upload_base64_to_r2(data.avatar, "avatars", username)
        updates.append("avatar = ?")
        params.append(avatar_url)
        
    if updates:
        query = f"UPDATE users SET {', '.join(updates)} WHERE id = ?"
        params.append(user_id)
        await execute_d1_query(query, params)
        
    # Fetch updated user
    updated_users = await execute_d1_query("SELECT * FROM users WHERE id = ?", [user_id])
    updated_user = updated_users[0]
    
    return {
        "success": True,
        "message": "Profile updated successfully",
        "user": {
            "id": updated_user["id"],
            "email": updated_user["email"],
            "username": updated_user.get("username") or updated_user["email"].split('@')[0],
            "avatar": updated_user.get("avatar")
        }
    }