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"Universal Captcha Solver 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" # သို့မဟုတ် gemini-2.0-flash-exp

client = genai.Client(api_key=GEMINI_API_KEY)

# --- Main Prompt ---
# ဒီ Prompt က Challenge အမျိုးအစားအားလုံးကို ခွဲခြားနိုင်အောင် ရေးထားတာဖြစ်ပါတယ်
UNIVERSAL_PROMPT = """
Analyze this hCaptcha challenge image. Perform the specific task requested in the image instruction.

Task Types and Requirements:
1. "Drag the blocks/letter": Identify the source block and the exact target slot. 
   Return: {"type": "drag", "moves": [{"source": [x,y], "destination": [x,y]}]}

2. "Click/Pick/Select objects" (Ingredients, Wall outlets, Bicycles, etc.): Identify all matching items.
   Return: {"type": "click", "points": [[x,y], [x,y], ...]}

3. "Two shapes that are identical": Find the two matching shapes.
   Return: {"type": "click", "points": [[x,y], [x,y]]}

4. "Click the asterisk where animals meet": Identify the meeting point.
   Return: {"type": "click", "points": [[x,y]]}

Output Rules:
- Return ONLY a strict JSON object.
- Use normalized coordinates [0-1000] for both X and Y.
- Ensure points are the exact centers of the objects.
"""

# --- Helper: Coordinate Scaling ---
def scale_coord(p, target_w=500, target_h=470):
    try:
        # 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]
        return {
            "x": int((float(raw_x) / 1000) * target_w),
            "y": int((float(raw_y) / 1000) * target_h)
        }
    except:
        return None

async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    photo_file = await update.message.photo[-1].get_file()
    photo_path = f"task_{update.message.message_id}.png"
    await photo_file.download_to_drive(photo_path)
    
    await update.message.reply_text("🧠 Analyzing challenge type and calculating coordinates...")

    try:
        raw_img = Image.open(photo_path)
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=[UNIVERSAL_PROMPT, raw_img]
        )
        
        # JSON parsing logic
        raw_text = response.text
        match = re.search(r"\{.*\}", raw_text, re.DOTALL)
        data = json.loads(match.group(0)) if match else {}
        
        result = {"model": MODEL_ID, "analysis_type": data.get("type")}
        
        # Challenge အမျိုးအစားအလိုက် Scaling လုပ်ခြင်း
        if data.get("type") == "drag":
            moves = []
            for m in data.get("moves", []):
                moves.append({
                    "source": scale_coord(m["source"]),
                    "destination": scale_coord(m["destination"])
                })
            result["required_moves"] = moves
        else:
            points = [scale_coord(p) for p in data.get("points", []) if scale_coord(p)]
            result["points_to_click"] = points

        await update.message.reply_text(f"✅ Solution:\n\n```json\n{json.dumps(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 started. Send any hCaptcha image!")
    app.run_polling()