from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import base64
import cv2
import numpy as np
import random
import uvicorn

app = FastAPI()

# 啟用 CORS 支援（允許 Unity WebGL 或跨網域連線）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 定義 Unity 傳過來的資料格式
class ClientRequest(BaseModel):
    message: str

# 輔助函式：動態生成一張隨機的 2D 圖片，並轉為 Base64 字串
def generate_random_image_base64():
    # 1. 建立一張 400x400 的空白黑色圖片 (3 個顏色通道)
    width, height = 400, 400
    img = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 2. 隨機填滿背景顏色
    bg_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
    img[:] = bg_color
    
    # 3. 在圖片上畫一些隨機的彩色圓形或方形
    for _ in range(5):
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        thickness = random.choice([-1, 2, 4]) # -1 代表填滿
        
        shape_type = random.choice(["circle", "rectangle"])
        if shape_type == "circle":
            center = (random.randint(0, width), random.randint(0, height))
            radius = random.randint(20, 80)
            cv2.circle(img, center, radius, color, thickness)
        else:
            pt1 = (random.randint(0, width), random.randint(0, height))
            pt2 = (random.randint(0, width), random.randint(0, height))
            cv2.rectangle(img, pt1, pt2, color, thickness)

    # 4. 將 OpenCV 的圖片物件編碼為 JPG 格式的二進位數據
    _, buffer = cv2.imencode('.jpg', img)
    
    # 5. 將二進位數據轉為 Base64 字串
    img_base64 = base64.b64encode(buffer).decode('utf-8')
    return img_base64
    
@app.get("/")
def read_root():
    return {"message": "FastAPI server is running successfully on HF Spaces!"}
    
@app.post("/get_image_and_text")
def handle_unity_request(request: ClientRequest):
    # 檢查客戶端傳來的是不是 "ok" (不區分大小寫)
    if request.message.strip().lower() != "ok":
        raise HTTPException(status_code=400, detail="Invalid message. Please send 'ok'.")
    
    # 生成隨機圖片的 Base64 字串
    base64_image = generate_random_image_base64()
    
    # 準備隨機的回覆文字
    random_responses = [
        "Server received OK! Here is your lucky image.",
        "System Status: Perfect. Image generated successfully.",
        "Hello from Hugging Face! Processing done.",
        "OK received. Sending custom 2D matrix back to Unity."
    ]
    selected_text = random.choice(random_responses)
    
    # 回傳給 Unity
    return {
        "status": "success",
        "reply_text": selected_text,
        "image_base64": base64_image
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)