import os
import json
import asyncio
import threading
import logging
from http.server import HTTPServer, BaseHTTPRequestHandler
from PIL import Image
from google import genai
from telegram import Update
from telegram.ext import ApplicationBuilder, ContextTypes, MessageHandler, filters

# --- LOGGING ---
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)

# --- PORT 7860 HEALTH CHECK (For Hugging Face) ---
class HealthCheckHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b"Gemini 3.1 Flash Lite is Online!")

def run_health_check():
    server = HTTPServer(('0.0.0.0', 7860), HealthCheckHandler)
    server.serve_forever()

# --- CONFIGURATION ---
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
TELEGRAM_TOKEN = os.getenv("TELEGRAM_TOKEN")
MODEL_ID = "gemini-3.1-flash-lite-preview" 

client = genai.Client(api_key=GEMINI_API_KEY)

async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    photo_file = await update.message.photo[-1].get_file()
    photo_path = "hcaptcha_input.png"
    await photo_file.download_to_drive(photo_path)
    
    await update.message.reply_text("Gemini 3.1 Flash Lite (Normalized Mode) ဖြင့် စစ်ဆေးနေပါပြီ...")

    # AI ကို 0-1000 scale နဲ့ပဲ တွက်ခိုင်းတဲ့ Prompt
    prompt = """
    Analyze this hCaptcha image.
    Task: Find center coordinates of objects requiring a wall outlet (Lamps/Monitors).

    Spatial Logic:
    - Target 1: Square icon on the LEFT (X < 200).
    - Target 2: Pentagon icon in the CENTER (X ≈ 275).

    Output Rules:
    1. Return ONLY strict JSON.
    2. Use normalized coordinates [0-1000] for both X and Y.
       (0 is top/left, 1000 is bottom/right).
    3. Output Format:
    {
      "detected_items": ["item1", "item2"],
      "solution_normalized": [[x1, y1], [x2, y2]]
    }
    """

    try:
        raw_img = Image.open(photo_path)
        
        # Gemini API Call
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=[prompt, raw_img],
            config={'response_mime_type': 'application/json'}
        )
        
        # JSON Parsing & Scaling Logic
        data = json.loads(response.text.strip())
        normalized_points = data.get("solution_normalized", [])
        
        final_solution = []
        for p in normalized_points:
            # Formula: Pixel = (Normalized / 1000) * Max_Pixel
            pixel_x = int((p[0] / 1000) * 500)
            pixel_y = int((p[1] / 1000) * 470)
            final_solution.append({"point": [pixel_x, pixel_y]})
            
        final_result = {
            "model": "Gemini 3.1 Flash Lite",
            "detected_items": data.get("detected_items"),
            "solution": final_solution
        }
        
        # Telegram ဆီ အဖြေပြန်ပို့ခြင်း
        await update.message.reply_text(f"✅ Scaled Result (500x470):\n\n{json.dumps(final_result, indent=2)}")
        
    except Exception as e:
        await update.message.reply_text(f"❌ Error: {str(e)}")
    finally:
        if os.path.exists(photo_path):
            os.remove(photo_path)

if __name__ == '__main__':
    # 1. Start Web Server
    threading.Thread(target=run_health_check, daemon=True).start()
    
    # 2. Start Telegram Bot
    app = ApplicationBuilder().token(TELEGRAM_TOKEN).build()
    app.add_handler(MessageHandler(filters.PHOTO, handle_photo))
    
    print("Bot is ready and using Normalized Scaling logic!")
    app.run_polling()