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)

# --- Hugging Face Health Check Server ---
class HealthCheckHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b"Coordinate Finder Bot (Drag-and-Drop 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!")

# Use genai.Client() for compatibility with older library versions on Hugging Face
client = genai.Client(api_key=GEMINI_API_KEY)


# --- Telegram Bot's Main Logic ---
async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """This function is called when a user sends a photo to the bot."""
    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 များကို ရှာဖွေနေပါသည်...")

    # --- This prompt is specifically designed for drag-and-drop challenges ---
    prompt = """
    Analyze this hCaptcha drag-and-drop challenge image.

    Task:
    1. Find the center coordinate of the source object to be dragged.
    2. Find the center coordinates of all possible target destination images.

    Output Rules:
    1. Return ONLY a single, strict JSON object. Do not add any extra text or markdown like ```json.
    2. Use normalized coordinates [0-1000] for both X and Y axes.
    3. The JSON must have "source_normalized" (a single point) and "destinations_normalized" (a list of points).

    Example JSON Output Format:
    {
      "source_normalized": [x_start, y_start],
      "destinations_normalized": [
        [x_end_1, y_end_1],
        [x_end_2, y_end_2]
      ]
    }
    """

    try:
        raw_img = Image.open(photo_path)
        
        # Call the Gemini API
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=[prompt, raw_img]
        )
        
        # Clean the AI's response to ensure it's valid JSON
        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)
        
        # --- Parse and scale the coordinates ---
        source_norm = data.get("source_normalized",)
        source_pixel = {
            "x": int((source_norm / 1000) * 500),
            "y": int((source_norm / 1000) * 470)
        }

        destinations_norm = data.get("destinations_normalized", [])
        destinations_pixel = []
        for p in destinations_norm:
            pixel_point = {
                "x": int((p / 1000) * 500),
                "y": int((p / 1000) * 470)
            }
            destinations_pixel.append(pixel_point)
            
        final_result = {
            "model": MODEL_ID,
            "analysis_type": "drag-and-drop",
            "source_location": source_pixel,
            "destination_locations": destinations_pixel
        }
        
        await update.message.reply_text(f"✅ Coordinate Finder Result:\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 Coordinate Finder Bot...")
    app = ApplicationBuilder().token(TELEGRAM_TOKEN).build()
    app.add_handler(MessageHandler(filters.PHOTO, handle_photo))
    
    logging.info("Bot is ready. Send an image to find coordinates.")
    app.run_polling()