import os
import json
import re
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 Setup ---
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)

# --- Health Check Server ---
class HealthCheckHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b"AI Captcha Solver Fixed is Online!")

def run_health_check():
    httpd = HTTPServer(('0.0.0.0', 7860), HealthCheckHandler)
    logging.info("Health check server started on port 7860.")
    httpd.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)

# --- Core Logic with Error Fix ---
async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    photo_file = await update.message.photo[-1].get_file()
    photo_path = f"input_{update.message.message_id}.png"
    await photo_file.download_to_drive(photo_path)
    
    await update.message.reply_text(f"🔎 {MODEL_ID} ဖြင့် Coordinate များ တွက်ချက်နေပါသည်...")

    prompt = """
    Analyze this hCaptcha image. 
    1. Identify all objects matching the instructions.
    2. Return their center coordinates as a list of points.
    
    Output Format (Strict JSON):
    {
      "selected_points_normalized": [[x1, y1], [x2, y2]]
    }
    """

    try:
        raw_img = Image.open(photo_path)
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=[prompt, raw_img]
        )
        
        # JSON Cleaning
        raw_text = response.text
        match = re.search(r"\{.*\}", raw_text, re.DOTALL)
        clean_json_text = match.group(0) if match else raw_text.strip()
        data = json.loads(clean_json_text)
        
        selected_points = data.get("selected_points_normalized", [])
        click_locations = []

        # --- THE FIX: Robust Coordinate Scaling ---
        for p in selected_points:
            try:
                # p သည် [x, y] ဖြစ်နေကြောင်း သေချာအောင် စစ်ဆေးခြင်း
                if isinstance(p, list) and len(p) >= 2:
                    # အကယ်၍ p[0] သို့မဟုတ် p[1] က list ထပ်ဖြစ်နေလျှင် (Nested list error fix)
                    raw_x = p[0][0] if isinstance(p[0], list) else p[0]
                    raw_y = p[1][0] if isinstance(p[1], list) else p[1]
                    
                    # ကိန်းဂဏန်းအဖြစ် ပြောင်းလဲခြင်း
                    val_x = float(raw_x)
                    val_y = float(raw_y)
                    
                    point = {
                        "x": int((val_x / 1000) * 500),
                        "y": int((val_y / 1000) * 470)
                    }
                    click_locations.append(point)
            except (TypeError, ValueError, IndexError):
                continue # Data format လွဲနေသော point ကို ကျော်သွားမည်

        final_result = {
            "model": MODEL_ID,
            "analysis_type": "click-selection",
            "points_to_click": click_locations
        }
        
        await update.message.reply_text(
            f"✅ **Fixed Solution:**\n\n```json\n{json.dumps(final_result, indent=2)}\n```",
            parse_mode='Markdown'
        )
        
    except Exception as e:
        logging.error(f"Error: {e}")
        await update.message.reply_text(f"❌ Error: {str(e)}")
    finally:
        if os.path.exists(photo_path):
            os.remove(photo_path)

if __name__ == '__main__':
    threading.Thread(target=run_health_check, daemon=True).start()
    app = ApplicationBuilder().token(TELEGRAM_TOKEN).build()
    app.add_handler(MessageHandler(filters.PHOTO, handle_photo))
    logging.info("Bot is running...")
    app.run_polling()