app.py

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 """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <title>AI Puzzle Solver Dashboard</title>
        <style>
            body { background: #0d1117; color: #c9d1d9; font-family: 'Segoe UI', sans-serif; margin: 20px; }
            .container { display: grid; grid-template-columns: 1fr 1fr; gap: 20px; max-width: 1200px; margin: auto; }
            .box { background: #161b22; border: 1px solid #30363d; border-radius: 8px; padding: 20px; min-height: 500px; }
            h2 { color: #58a6ff; font-size: 14px; margin-top: 0; text-transform: uppercase; border-bottom: 1px solid #30363d; padding-bottom: 10px; }
            #logs { font-family: 'Courier New', monospace; font-size: 12px; height: 450px; overflow-y: auto; color: #8b949e; }
            .log-entry { margin-bottom: 5px; border-bottom: 1px solid #21262d; padding-bottom: 5px; }
            .log-success { color: #3fb950; font-weight: bold; }
            .log-error { color: #f85149; font-weight: bold; }
            .log-warn { color: #d29922; font-weight: bold; }
            #vision-img { width: 100%; border-radius: 4px; border: 2px solid #30363d; background: #000; display: none; }
            .placeholder { text-align: center; color: #8b949e; padding-top: 100px; font-style: italic; }
        </style>
    </head>
    <body>
        <h1 style="font-size: 22px; text-align: center; color: #f0f6fc; margin-bottom: 30px;">⚡ AI Block Puzzle Solver Dashboard</h1>
        <div class="container">
            <div class="box">
                <h2>System Logs</h2>
                <div id="logs">Connecting to AI...</div>
            </div>
            <div class="box">
                <h2>Latest Vision (Move Lines)</h2>
                <div id="placeholder" class="placeholder">Waiting for image...</div>
                <img id="vision-img" src="" alt="Latest Result">
                <p style="font-size: 11px; color: #8b949e; text-align: center; margin-top: 15px;">Auto-refresh active (2s)</p>
            </div>
        </div>
        <script>
            async function update() {
                try {
                    const lRes = await fetch('/get_logs');
                    const lData = await lRes.json();
                    document.getElementById('logs').innerHTML = lData.logs.map(m => {
                        let c = m.includes('SUCCESS') ? 'log-success' : 
                               (m.includes('ERROR') ? 'log-error' : 
                               (m.includes('WARNING') ? 'log-warn' : ''));
                        return `<div class="log-entry ${c}">${m}</div>`;
                    }).join('');

                    const i = document.getElementById('vision-img');
                    const p = document.getElementById('placeholder');
                    const iRes = await fetch('/get_latest_vision');
                    if (iRes.ok) {
                        i.src = '/get_latest_vision?t=' + new Date().getTime();
                        i.style.display = 'block'; p.style.display = 'none';
                    }
                } catch (e) {}
            }
            setInterval(update, 2000);
        </script>
    </body>
    </html>
    """

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