import os
import time
import json
import re
import io
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image, ImageDraw
import google.generativeai as genai
# --- CONFIGURATION ---
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
if GEMINI_API_KEY:
genai.configure(api_key=GEMINI_API_KEY)
# Model
model = genai.GenerativeModel('gemini-3.1-flash-lite-preview')
app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
LOGS = []
LATEST_IMAGE_PATH = "/tmp/latest_vision.png"
def add_log(message):
timestamp = time.strftime("%H:%M:%S")
LOGS.append(f"[{timestamp}] {message}")
if len(LOGS) > 20: LOGS.pop(0)
def draw_red_lines(image_bytes, moves):
""" Start ကနေ End ကို အနီရောင်မျဉ်းဆွဲပေးမယ့် function """
try:
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
draw = ImageDraw.Draw(img)
width, height = img.size
count = 0
for move in moves:
start = move.get("start", {})
end = move.get("end", {})
# Start and End coordinates များကို ယူခြင်း
sx, sy = start.get('x', 0), start.get('y', 0)
ex, ey = end.get('x', 0), end.get('y', 0)
# Data Type Check (List ဖြစ်နေရင် ပထမတန်ဖိုးကို ယူမယ်)
sx = sx[0] if isinstance(sx, list) else sx
sy = sy[0] if isinstance(sy, list) else sy
ex = ex[0] if isinstance(ex, list) else ex
ey = ey[0] if isinstance(ey, list) else ey
# ကိန်းဂဏန်းတွေ ဖြစ်မှ မျဉ်းဆွဲမယ်
if all(isinstance(v, (int, float)) for v in [sx, sy, ex, ey]):
# Normalized (0-1000) ကို Pixel ပြောင်းခြင်း
px_start = ((sx / 1000) * width, (sy / 1000) * height)
px_end = ((ex / 1000) * width, (ey / 1000) * height)
# ၁။ မျဉ်းနီဆွဲခြင်း (Red Line)
draw.line([px_start, px_end], fill="red", width=6)
# ၂။ စတင်ရမယ့်နေရာကို အပြာစက်လေး ပြခြင်း (Start Point)
r = 10
draw.ellipse([px_start[0]-r, px_start[1]-r, px_start[0]+r, px_start[1]+r], fill="blue", outline="white", width=2)
# ၃။ သွားချရမယ့်နေရာကို အစိမ်းစက်လေး ပြခြင်း (End Point)
draw.ellipse([px_end[0]-r, px_end[1]-r, px_end[0]+r, px_end[1]+r], fill="green", outline="white", width=2)
count += 1
img.save(LATEST_IMAGE_PATH)
return count
except Exception as e:
add_log(f"Draw Logic Error: {str(e)}")
return 0
@app.post("/solve/image")
async def solve_image(file: UploadFile = File(...)):
add_log(f"Received Image: {file.filename}")
img_content = await file.read()
# 💥 Drag & Drop အတွက် ပြောင်းလဲထားသော Prompt
prompt = """
You are an AI solving a block puzzle CAPTCHA.
The user needs to drag the blocks on the right side and drop them onto the matching icon patterns on the grid on the left.
Find the center coordinate of each block (start) and the center coordinate of where it should go on the grid (end).
Return ONLY a valid JSON object.
Use 0-1000 normalized coordinates.
Format exactly like this:
{"moves": [
{"start": {"x": 800, "y": 300}, "end": {"x": 350, "y": 450}}
]}
CRITICAL: "x" and "y" must be single numbers, NOT lists. No extra text.
"""
try:
response = model.generate_content([
prompt,
{"mime_type": "image/jpeg", "data": img_content}
])
raw_text = response.text.strip()
# Clean up markdown format
raw_text = re.sub(r'```[a-z]*\n?|```', '', raw_text, flags=re.IGNORECASE).strip()
json_match = re.search(r'\{.*\}', raw_text, re.DOTALL)
if json_match:
json_str = json_match.group().replace("'", '"')
try:
data = json.loads(json_str)
except:
# Fallback clean-up
json_str = re.sub(r',\s*}', '}', json_str)
json_str = re.sub(r',\s*]', ']', json_str)
data = json.loads(json_str)
moves = data.get("moves", [])
processed_count = draw_red_lines(img_content, moves)
if processed_count > 0:
add_log(f"SUCCESS: Drawn {processed_count} move lines.")
return {"success": True, "moves": moves}
else:
add_log("WARNING: No valid moves to draw.")
return {"success": False, "error": "No moves extracted"}
else:
add_log("ERROR: AI failed to return JSON.")
return {"success": False, "error": "No JSON found"}
except Exception as e:
add_log(f"SYSTEM ERROR: {str(e)}")
return {"success": False, "error": str(e)}
@app.get("/get_latest_vision")
async def get_latest_vision():
if os.path.exists(LATEST_IMAGE_PATH):
return FileResponse(LATEST_IMAGE_PATH)
return HTMLResponse(status_code=404)
@app.get("/get_logs")
async def get_logs():
return {"logs": LOGS}
@app.get("/", response_class=HTMLResponse)
async def dashboard():
return """
AI Puzzle Solver Dashboard
⚡ AI Block Puzzle Solver Dashboard
System Logs
Connecting to AI...
Latest Vision (Move Lines)
Waiting for image...
![Latest Result]()
Auto-refresh active (2s)
"""
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)