Update server.py

server.py

@@ -1,88 +1,147 @@
+"""
+Backend API HỢP NHẤT (Web + Desktop) - HT MATH V6
+Chạy trên Hugging Face Spaces (Docker Version)
+"""
+
 import os
-import random
-import logging
+import io
+import time
+import asyncio
+import re
+import tempfile
+import hashlib
+import secrets
+import uuid
+import math
 import base64
-from io import BytesIO
-from typing import Optional, List
+import random
+from typing import List, Optional
 
-from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Body
+# --- THƯ VIỆN CHÍNH ---
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Request, Body
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import google.generativeai as genai
-from dotenv import load_dotenv
+from fastapi.responses import JSONResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel # Thêm Pydantic cho Desktop App
+
+# --- THƯ VIỆN XỬ LÝ ẢNH & AI ---
 from PIL import Image
+import fitz  # PyMuPDF
+import google.generativeai as genai
+
+# --- PANDOC IMPORT ---
+try:
+    import pypandoc
+    print(f"INFO: Pandoc version detected: {pypandoc.get_pandoc_version()}")
+except ImportError:
+    print("CRITICAL WARNING: pypandoc module not found.")
+except OSError:
+    print("CRITICAL WARNING: pandoc binary not found in system path.")
+
+# --- SUPABASE ---
+try:
+    from supabase import create_client, Client
+    SUPABASE_AVAILABLE = True
+except ImportError:
+    SUPABASE_AVAILABLE = False
+    Client = None
+    create_client = None
 
-# --- LOAD BIẾN MÔI TRƯỜNG ---
-load_dotenv()
+# ===== CẤU HÌNH =====
+# Load biến môi trường
+GEMINI_API_KEYS = os.getenv("GEMINI_API_KEYS", "").split(",")
+# Lọc bỏ key rỗng
+GEMINI_API_KEYS = [k.strip() for k in GEMINI_API_KEYS if k.strip()]
 
-# --- CẤU HÌNH SERVER ---
+GEMINI_MODELS = os.getenv("GEMINI_MODELS", "gemini-2.5-flash,gemini-1.5-pro").split(",")
+GEMINI_MODELS = [m.strip() for m in GEMINI_MODELS if m.strip()]
+
+SUPABASE_URL = os.getenv("SUPABASE_URL")
+SUPABASE_KEY = os.getenv("SUPABASE_KEY")
+
+# Khởi tạo Supabase
+supabase: Optional[Client] = None
+if SUPABASE_AVAILABLE and SUPABASE_URL and SUPABASE_KEY:
+    try:
+        supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
+        print("INFO: Supabase connected successfully.")
+    except Exception as e:
+        print(f"ERROR: Failed to connect to Supabase: {e}")
+
+# Khởi tạo FastAPI
 app = FastAPI(title="HT MATH UNIFIED SERVER")
 
-# Cấu hình CORS (Cho phép cả Web và Desktop App gọi vào)
+# Cấu hình CORS (Cho phép cả Web và Desktop App)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Cho phép tất cả các nguồn (Web, App Desktop)
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# --- PHẦN 1: CẤU HÌNH CHO DESKTOP APP (Key Rotation & Models) ---
+# --- HELPER FUNCTIONS ---
+def get_random_api_key():
+    if not GEMINI_API_KEYS:
+        raise HTTPException(status_code=500, detail="Server chưa cấu hình GEMINI_API_KEYS")
+    return secrets.choice(GEMINI_API_KEYS)
 
-# 1. Lấy danh sách API KEYS từ biến môi trường
-# Format trên HuggingFace: Key1,Key2,Key3
-api_keys_env = os.getenv("GEMINI_API_KEYS", "")
-DESKTOP_API_KEYS = [k.strip() for k in api_keys_env.split(",") if k.strip()]
+def clean_latex_formulas(text):
+    """Làm sạch và chuẩn hóa LaTeX"""
+    text = re.sub(r'\\\(', '$', text)
+    text = re.sub(r'\\\)', '$', text)
+    text = re.sub(r'\\\[', '$$', text)
+    text = re.sub(r'\\\]', '$$', text)
+    return text
 
-# 2. Lấy danh sách MODELS
-models_env = os.getenv("GEMINI_MODELS", "gemini-2.0-flash-exp,gemini-1.5-pro,gemini-1.5-flash")
-DESKTOP_MODELS = [m.strip() for m in models_env.split(",") if m.strip()]
+# ==============================================================================
+# PHẦN 1: API DÀNH RIÊNG CHO DESKTOP APP (HT MATH V6 CLIENT)
+# ==============================================================================
 
-# Model dữ liệu nhận từ Desktop App
 class DesktopGenerateRequest(BaseModel):
     prompt: str
     model: Optional[str] = "gemini-1.5-flash"
     image: Optional[str] = None # Base64 string
 
-@app.get("/")
-async def root():
-    return {
-        "status": "online",
-        "server": "HT MATH UNIFIED SERVER",
-        "desktop_keys_loaded": len(DESKTOP_API_KEYS),
-        "web_support": "Active"
-    }
-
-# --- API 1: Lấy danh sách Model (Cho Desktop App) ---
 @app.get("/api/models")
 async def get_models_desktop():
-    return {"models": DESKTOP_MODELS}
+    """API trả về danh sách model cho Desktop App cập nhật vào ComboBox"""
+    # Nếu danh sách rỗng, trả về default để tránh lỗi app
+    models = GEMINI_MODELS if GEMINI_MODELS else ["gemini-1.5-flash"]
+    return {"models": models}
 
-# --- API 2: Xử lý AI (Cho Desktop App - Có xoay vòng Key) ---
 @app.post("/api/generate")
 async def generate_content_desktop(req: DesktopGenerateRequest):
-    if not DESKTOP_API_KEYS:
-        raise HTTPException(status_code=500, detail="Server chưa cấu hình GEMINI_API_KEYS")
-
+    """
+    API xử lý AI cho Desktop App.
+    Khác với Web App (nhận Multipart), API này nhận JSON chứa Base64 image.
+    """
     try:
-        # 1. Chọn ngẫu nhiên 1 Key (Load Balancing)
-        current_key = random.choice(DESKTOP_API_KEYS)
-        genai.configure(api_key=current_key)
+        # 1. Chọn Key ngẫu nhiên (Load Balancing)
+        api_key = get_random_api_key()
+        genai.configure(api_key=api_key)
         
         # 2. Chọn Model
-        # Nếu model client gửi lên không có trong danh sách hỗ trợ, dùng model mặc định
-        model_name = req.model if req.model in DESKTOP_MODELS else DESKTOP_MODELS[0]
+        # Nếu model client gửi lên không có trong danh sách hỗ trợ, dùng model đầu tiên
+        model_name = req.model
+        if model_name not in GEMINI_MODELS and GEMINI_MODELS:
+            model_name = GEMINI_MODELS[0]
+            
         model = genai.GenerativeModel(model_name)
 
         # 3. Chuẩn bị nội dung gửi đi
         content_parts = [req.prompt]
 
-        # 4. Xử lý ảnh (nếu có)
+        # 4. Xử lý ảnh (Base64 -> Image)
         if req.image:
             try:
                 # Desktop App gửi ảnh dạng Base64 string
+                # Cần xử lý trường hợp có prefix data:image/...;base64,
+                if "," in req.image:
+                    req.image = req.image.split(",")[1]
+                
                 image_bytes = base64.b64decode(req.image)
-                image = Image.open(BytesIO(image_bytes))
+                image = Image.open(io.BytesIO(image_bytes))
                 content_parts.append(image)
             except Exception as e:
                 raise HTTPException(status_code=400, detail=f"Lỗi xử lý ảnh base64: {str(e)}")
@@ -96,27 +155,194 @@ async def generate_content_desktop(req: DesktopGenerateRequest):
             raise HTTPException(status_code=500, detail="Gemini không trả về nội dung text.")
 
     except Exception as e:
-        print(f"Lỗi Desktop API: {e}")
+        print(f"Error Desktop API: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 
 # ==============================================================================
-# --- PHẦN 2: KHU VỰC CỦA WEB APP (COPY CODE CŨ VÀO DƯỚI ĐÂY) ---
+# PHẦN 2: API DÀNH CHO WEB APP (HT MATH WEB V6)
 # ==============================================================================
 
-# Bạn hãy dán các hàm xử lý của Web App vào đây.
-# Ví dụ: import supabase, các hàm convert_file, login, signup...
-# Đảm bảo giữ nguyên logic cũ của Web App.
+# --- AUTHENTICATION (WEB) ---
+@app.post("/api/auth/register")
+async def register(request: Request):
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database service unavailable")
+    
+    data = await request.json()
+    email = data.get("email")
+    password = data.get("password")
+    full_name = data.get("full_name")
+
+    if not email or not password:
+        raise HTTPException(status_code=400, detail="Vui lòng nhập Email và Mật khẩu")
+
+    try:
+        # 1. Đăng ký Auth user
+        auth_res = supabase.auth.sign_up({
+            "email": email,
+            "password": password,
+            "options": {"data": {"full_name": full_name}}
+        })
+
+        if not auth_res.user:
+             raise HTTPException(status_code=400, detail="Đăng ký thất bại (Auth)")
+
+        # 2. Lưu vào bảng users (public)
+        user_data = {
+            "id": auth_res.user.id,
+            "email": email,
+            "full_name": full_name,
+            "role": "user",
+            "created_at": "now()"
+        }
+        supabase.table("users").insert(user_data).execute()
+
+        return {"success": True, "message": "Đăng ký thành công! Vui lòng kiểm tra email xác nhận."}
+
+    except Exception as e:
+        print(f"Register Error: {str(e)}")
+        # Xử lý lỗi Supabase trả về
+        msg = str(e)
+        if "User already registered" in msg:
+             raise HTTPException(status_code=400, detail="Email này đã được đăng ký.")
+        raise HTTPException(status_code=500, detail=f"Lỗi đăng ký: {msg}")
+
+@app.post("/api/auth/login")
+async def login(request: Request):
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database service unavailable")
+
+    data = await request.json()
+    email = data.get("email")
+    password = data.get("password")
+
+    try:
+        res = supabase.auth.sign_in_with_password({"email": email, "password": password})
+        if res.user:
+            # Lấy thông tin role từ bảng users
+            user_info = supabase.table("users").select("*").eq("id", res.user.id).execute()
+            role = "user"
+            full_name = ""
+            if user_info.data:
+                role = user_info.data[0].get("role", "user")
+                full_name = user_info.data[0].get("full_name", "")
+
+            return {
+                "success": True,
+                "access_token": res.session.access_token,
+                "user": {
+                    "id": res.user.id,
+                    "email": res.user.email,
+                    "role": role,
+                    "full_name": full_name
+                }
+            }
+        raise HTTPException(status_code=401, detail="Email hoặc mật khẩu không đúng")
+    except Exception as e:
+        raise HTTPException(status_code=401, detail=str(e))
+
+# --- IMAGE PROCESSING UTILS (WEB) ---
+async def process_large_image(image: Image.Image, model_name: str, prompt: str, semaphore: asyncio.Semaphore) -> str:
+    """Xử lý ảnh lớn bằng cách cắt nhỏ (Overlap Stitching)"""
+    width, height = image.size
+    
+    # Nếu ảnh nhỏ, xử lý trực tiếp
+    if height < 2000:
+        async with semaphore:
+             return await call_gemini_vision(image, model_name, prompt)
 
-# Ví dụ mẫu (Nếu bạn dùng Supabase):
-# from supabase import create_client, Client
-# url: str = os.environ.get("SUPABASE_URL")
-# key: str = os.environ.get("SUPABASE_KEY")
-# supabase: Client = create_client(url, key)
+    # Cấu hình cắt ảnh
+    segment_height = 1500
+    overlap = 300
+    segments = []
+    
+    for y in range(0, height, segment_height - overlap):
+        box = (0, y, width, min(y + segment_height, height))
+        segment = image.crop(box)
+        segments.append(segment)
+        if y + segment_height >= height:
+            break
 
-# @app.post("/convert")
-# async def convert_file(...):
-#     ... code cũ của bạn ...
+    # Gọi API song song cho các phần
+    tasks = []
+    for seg in segments:
+        tasks.append(call_gemini_vision(seg, model_name, prompt))
+    
+    results = await asyncio.gather(*tasks)
+    return "\n".join(results) # Ghép kết quả đơn giản
+
+async def call_gemini_vision(image: Image.Image, model_name: str, prompt: str) -> str:
+    """Hàm wrapper gọi Gemini Vision"""
+    try:
+        api_key = get_random_api_key()
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel(model_name)
+        response = await model.generate_content_async([prompt, image])
+        return response.text if response.text else ""
+    except Exception as e:
+        print(f"Gemini Error: {e}")
+        return ""
+
+# --- MAIN API: PROCESS IMAGE (WEB) ---
+@app.post("/api/process-image")
+async def process_image_web(
+    file: UploadFile = File(...),
+    model: str = Form("gemini-1.5-pro"),
+    prompt: str = Form("Hãy chuyển đổi nội dung trong ảnh thành định dạng Markdown LaTeX.")
+):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+        
+        # Giới hạn số luồng xử lý đồng thời để tránh Rate Limit
+        global_semaphore = asyncio.Semaphore(5) 
+        
+        # Xử lý ảnh
+        result_text = await process_large_image(image, model, prompt, global_semaphore)
+        
+        return {"success": True, "result": clean_latex_formulas(result_text)}
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+# --- WORD EXPORT API (PANDOC NATIVE) ---
+@app.post("/api/export-docx")
+async def export_docx(markdown_text: str = Form(...)):
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".docx", delete=False) as tmp_file:
+            output_filename = tmp_file.name
+
+        # Dùng Pypandoc để convert
+        pypandoc.convert_text(
+            markdown_text,
+            to='docx',
+            format='markdown',
+            outputfile=output_filename,
+            extra_args=['--standalone']
+        )
+
+        return FileResponse(
+            output_filename,
+            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+            filename="Ket_qua_HT_MATH.docx"
+        )
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Lỗi xuất Word: {str(e)}")
+
+# --- TEST ENDPOINT ---
+@app.get("/")
+def home():
+    return {
+        "server": "HT MATH UNIFIED (Web + Desktop)",
+        "status": "online", 
+        "pandoc": "detected" if 'pypandoc' in globals() else "missing",
+        "desktop_api_ready": True
+    }
 
 if __name__ == "__main__":
     import uvicorn