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 ---
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"Gemini (Multi-Drag Mode) 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"

if not GEMINI_API_KEY or not TELEGRAM_TOKEN:
    logging.error("FATAL: GEMINI_API_KEY or TELEGRAM_TOKEN secret is not set!")

client = genai.Client(api_key=GEMINI_API_KEY)


# --- TELEGRAM BOT LOGIC ---
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} ဖြင့် Multi-Drag ပုံကိုစစ်ဆေးနေပါပြီ...")

    # !!!!! NEW PROMPT FOR MULTI-DRAG CHALLENGE !!!!!
    prompt = """
    Analyze this complex hCaptcha "Cover ALL icon patterns" challenge.

    Task:
    The user must drag multiple source blocks (purple, green, etc.) to cover all corresponding empty icon patterns in the central grid.
    Identify EVERY required drag-and-drop operation. For each operation, you must provide:
    1. The center coordinate of the source block to be dragged.
    2. The center coordinate of the destination empty pattern where the block should be dropped.

    Output Rules:
    1. Return ONLY a single, strict JSON object. Do not add any extra text or markdown.
    2. The JSON object must contain a single key: "drag_operations".
    3. "drag_operations" must be a LIST of objects.
    4. Each object in the list represents one drag action and must have two keys: "source_normalized" and "destination_normalized".
    5. Use normalized coordinates [0-1000] for all points.

    Example JSON Output Format:
    {
      "drag_operations": [
        { "source_normalized": [x1, y1], "destination_normalized": [x2, y2] },
        { "source_normalized": [x3, y3], "destination_normalized": [x4, y4] },
        { "source_normalized": [x5, y5], "destination_normalized": [x6, y6] }
      ]
    }
    """

    try:
        raw_img = Image.open(photo_path)
        
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=[prompt, raw_img]
        )
        
        # Clean potential markdown from the response
        raw_text = response.text
        match = re.search(r"```(json)?(.*)```", raw_text, re.DOTALL)
        clean_json_text = match.group(2).strip() if match else raw_text.strip()
            
        data = json.loads(clean_json_text)
        
        # --- NEW JSON Parsing & Scaling Logic for Multi-Drag ---
        drag_operations = data.get("drag_operations", [])
        scaled_operations = []
        
        for op in drag_operations:
            source_norm = op.get("source_normalized", [0, 0])
            dest_norm = op.get("destination_normalized", [0, 0])
            
            scaled_op = {
                "source_location": {
                    "x": int((source_norm[0] / 1000) * 500),
                    "y": int((source_norm[1] / 1000) * 470)
                },
                "destination_location": {
                    "x": int((dest_norm[0] / 1000) * 500),
                    "y": int((dest_norm[1] / 1000) * 470)
                }
            }
            scaled_operations.append(scaled_op)

        final_result = {
            "model": MODEL_ID,
            "analysis_type": "multi-drag-and-drop",
            "required_moves": scaled_operations
        }
        
        await update.message.reply_text(f"✅ Multi-Drag Solution (Scaled to 500x470):\n\n{json.dumps(final_result, indent=2)}")
        
    except Exception as e:
        logging.error(f"Error during processing: {e}")
        await update.message.reply_text(f"❌ Error: {str(e)}")
    finally:
        if os.path.exists(photo_path):
            os.remove(photo_path)

if __name__ == '__main__':
    logging.info("Starting health check server...")
    threading.Thread(target=run_health_check, daemon=True).start()
    
    logging.info("Starting Telegram bot...")
    app = ApplicationBuilder().token(TELEGRAM_TOKEN).build()
    app.add_handler(MessageHandler(filters.PHOTO, handle_photo))
    
    logging.info("Bot is ready and using Multi-Drag logic!")
    app.run_polling()