"""
hCaptcha Solver — Hugging Face Spaces
======================================
Required env secret:  GEMINI_API_KEY   (Google AI Studio key)
Optional env vars:
  GEMINI_MODEL       (default: gemini-2.0-flash)
  PROXY_SERVER       e.g. 1.2.3.4:5188
  PROXY_USERNAME
  PROXY_PASSWORD

Endpoints:
  GET  /                  → status check
  POST /solve/captcha     → solve hCaptcha
    body: { "url": "https://...", "sitekey": "..." }
    resp: { "success": true/false, "token": "P1_...", "error": "..." }
"""

import os
from contextlib import asynccontextmanager
from pathlib import Path
from typing import List

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

PORT           = int(os.environ.get("PORT", 7860))
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
GEMINI_MODEL   = os.environ.get("GEMINI_MODEL", "gemini-3.1-flash-lite-preview")

PROXY_SERVER   = os.environ.get("PROXY_SERVER",   "")
PROXY_USERNAME = os.environ.get("PROXY_USERNAME",  "")
PROXY_PASSWORD = os.environ.get("PROXY_PASSWORD",  "")

if not GEMINI_API_KEY:
    print("[WARNING] GEMINI_API_KEY is not set — solver will fail!")


# ── Patch GeminiProvider BEFORE any hcaptcha imports ────────────────────────

def _patch_gemini_provider():
    from google import genai as _genai
    from google.genai import types as _types
    from hcaptcha_challenger.tools.internal.providers.gemini import GeminiProvider

    MIME = {".jpg": "image/jpeg", ".jpeg": "image/jpeg",
            ".png": "image/png", ".webp": "image/webp"}

    @property  # type: ignore[misc]
    def _patched_client(self):
        if self._client is None:
            self._client = _genai.Client(api_key=GEMINI_API_KEY)
        return self._client

    GeminiProvider.client = _patched_client  # type: ignore[assignment]

    async def _patched_upload_files(self, files: List[Path]):
        parts = []
        for f in files:
            if f and Path(f).exists():
                data = Path(f).read_bytes()
                mime = MIME.get(Path(f).suffix.lower(), "image/png")
                parts.append(_types.Part.from_bytes(data=data, mime_type=mime))
        return parts

    GeminiProvider._upload_files = _patched_upload_files  # type: ignore[method-assign]

    @staticmethod  # type: ignore[misc]
    def _patched_files_to_parts(files):
        return files

    GeminiProvider._files_to_parts = _patched_files_to_parts  # type: ignore[method-assign]

    print("[patch] GeminiProvider patched → inline images")


_patch_gemini_provider()


# ── Patch RoboticArm to use cv2-first coordinate analysis ────────────────────

_cv_patch_applied = False

def _patch_robotic_arm_cv():
    global _cv_patch_applied
    if _cv_patch_applied:
        return
    _cv_patch_applied = True

    from contextlib import suppress
    from playwright.async_api import TimeoutError as PWTimeout
    from hcaptcha_challenger.agent.challenger import RoboticArm
    from hcaptcha_challenger.models import SpatialPath, PointCoordinate
    import cv_challenge as cv_mod

    async def _cv_challenge_image_drag_drop(self, job_type):
        frame_challenge = await self.get_challenge_frame_locator()
        crumb_count     = await self.check_crumb_count()
        cache_key       = self.config.create_cache_key(self.captcha_payload)

        challenge_view = frame_challenge.locator("//div[@class='challenge-view']")
        bbox = await challenge_view.bounding_box()

        for cid in range(crumb_count):
            await self.page.wait_for_timeout(
                self.config.WAIT_FOR_CHALLENGE_VIEW_TO_RENDER_MS
            )

            raw, projection = await self._capture_spatial_mapping(
                frame_challenge, cache_key, cid
            )
            user_prompt = self._match_user_prompt(job_type)

            try:
                import shutil as _sh
                _sh.copy(str(raw), f"/tmp/challenge_{job_type.value}_{cid}.png")
                if bbox:
                    _bstr = (f"x={bbox['x']:.0f} y={bbox['y']:.0f} "
                             f"w={bbox['width']:.0f} h={bbox['height']:.0f}")
                    print(f"[cv2][debug] screenshot → "
                          f"/tmp/challenge_{job_type.value}_{cid}.png  bbox={_bstr}")
            except Exception:
                pass

            cv_result = None
            if bbox:
                try:
                    cv_result = cv_mod.dispatch_by_type(
                        job_type.value, raw, bbox
                    )
                    if cv_result:
                        if isinstance(cv_result, list):
                            cv_mod.annotate(raw, cv_result, bbox,
                                            f"/tmp/cv2_annotated_{job_type.value}_{cid}.png")
                        print(f"[cv2][crumb {cid+1}] type={job_type.value} "
                              f"result={cv_result}")
                    else:
                        print(f"[cv2][crumb {cid+1}] type={job_type.value} "
                              f"→ no CV result, falling back to Gemini")
                except Exception as e:
                    print(f"[cv2][crumb {cid+1}] dispatch error: {e}")

            if cv_result is not None:
                if isinstance(cv_result, list):
                    for (fx, fy), (tx, ty) in cv_result:
                        path = SpatialPath(
                            start_point=PointCoordinate(x=int(fx), y=int(fy)),
                            end_point=PointCoordinate(x=int(tx), y=int(ty)),
                        )
                        await self._perform_drag_drop(path)

                elif isinstance(cv_result, tuple) and len(cv_result) == 2:
                    first = cv_result[0]
                    if isinstance(first, tuple):
                        (fx, fy), (tx, ty) = cv_result
                        path = SpatialPath(
                            start_point=PointCoordinate(x=int(fx), y=int(fy)),
                            end_point=PointCoordinate(x=int(tx), y=int(ty)),
                        )
                        await self._perform_drag_drop(path)
                    else:
                        cx, cy = cv_result
                        await self.page.mouse.click(int(cx), int(cy), delay=180)
                        await self.page.wait_for_timeout(500)

                with suppress(PWTimeout):
                    submit_btn = frame_challenge.locator(
                        "//div[@class='button-submit button']"
                    )
                    await self.click_by_mouse(submit_btn)

            else:
                print(f"[cv2][crumb {cid+1}] no CV result → using Gemini")
                response = await self._spatial_path_reasoner(
                    challenge_screenshot=raw,
                    grid_divisions=projection,
                    auxiliary_information=user_prompt,
                )
                import logging
                logging.getLogger(__name__).debug(
                    f"[{cid+1}/{crumb_count}]ToolInvokeMessage: {response.log_message}"
                )
                self._spatial_path_reasoner.cache_response(
                    path=cache_key.joinpath(f"{cache_key.name}_{cid}_model_answer.json")
                )
                for path in response.paths:
                    await self._perform_drag_drop(path)

                with suppress(PWTimeout):
                    submit_btn = frame_challenge.locator(
                        "//div[@class='button-submit button']"
                    )
                    await self.click_by_mouse(submit_btn)

    RoboticArm.challenge_image_drag_drop = _cv_challenge_image_drag_drop
    print("[patch] RoboticArm.challenge_image_drag_drop → cv2-first + Gemini fallback")


# ── FastAPI app ──────────────────────────────────────────────────────────────

class SolveRequest(BaseModel):
    url: str
    sitekey: str

class SolveResponse(BaseModel):
    success: bool
    token: str | None = None
    error: str | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    print("=" * 50)
    print("hCaptcha Solver — Hugging Face Space")
    print(f"Model  : {GEMINI_MODEL}")
    print(f"Port   : {PORT}")
    print(f"Key    : {'set' if GEMINI_API_KEY else 'NOT SET'}")
    print(f"Proxy  : {PROXY_SERVER}" if PROXY_SERVER else "Proxy  : none")
    print("=" * 50)
    yield


app = FastAPI(title="hCaptcha Solver", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])


async def run_solver(url: str, sitekey: str) -> dict | None:
    from playwright.async_api import async_playwright
    from hcaptcha_challenger.agent import AgentV, AgentConfig
    from hcaptcha_challenger.models import CaptchaResponse

    _patch_robotic_arm_cv()

    mock_html = f"""<!DOCTYPE html>
<html>
<head>
  <script src="https://js.hcaptcha.com/1/api.js" async defer></script>
  <style>body{{background:#111;display:flex;justify-content:center;align-items:center;height:100vh;margin:0}}</style>
</head>
<body>
  <div class="h-captcha" data-sitekey="{sitekey}"></div>
</body>
</html>"""

    proxy_config = None
    if PROXY_SERVER:
        proxy_config = {
            "server":   f"http://{PROXY_SERVER}",
            "username": PROXY_USERNAME,
            "password": PROXY_PASSWORD,
        }
        print(f"[proxy] Using {PROXY_SERVER}")

    async with async_playwright() as p:
        browser = await p.chromium.launch(
            headless=True,
            args=["--no-sandbox", "--disable-setuid-sandbox",
                  "--disable-dev-shm-usage", "--disable-gpu"],
            proxy=proxy_config,
        )
        context = await browser.new_context(
            viewport={"width": 1280, "height": 720},
            user_agent=(
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/122.0.0.0 Safari/537.36"
            ),
        )
        page = await context.new_page()

        await page.route(url, lambda route: route.fulfill(
            status=200, content_type="text/html", body=mock_html
        ))

        try:
            await page.goto(url, wait_until="load")
            await page.wait_for_selector(
                "//iframe[contains(@src, 'frame=checkbox')]", timeout=15000
            )

            agent_config = AgentConfig(
                GEMINI_API_KEY=GEMINI_API_KEY,
                CHALLENGE_CLASSIFIER_MODEL=GEMINI_MODEL,
                IMAGE_CLASSIFIER_MODEL=GEMINI_MODEL,
                SPATIAL_POINT_REASONER_MODEL=GEMINI_MODEL,
                SPATIAL_PATH_REASONER_MODEL=GEMINI_MODEL,
            )

            agent = AgentV(page=page, agent_config=agent_config)
            print(f"[*] Solving hCaptcha for {url}...")

            await agent.robotic_arm.click_checkbox()
            await agent.wait_for_challenge()

            if agent.cr_list:
                cr: CaptchaResponse = agent.cr_list[-1]
                return cr.model_dump(by_alias=True)

            return None
        finally:
            await browser.close()


@app.get("/")
async def root():
    return {
        "status": "online",
        "model": GEMINI_MODEL,
        "endpoints": {"POST /solve/captcha": "Solve hCaptcha"},
    }


@app.post("/solve/captcha", response_model=SolveResponse)
async def solve_captcha(request: SolveRequest):
    try:
        print("[solver] start")
        result = await run_solver(request.url, request.sitekey)
        print(f"[solver] done → pass={result.get('pass') if result else False}")
    except Exception as e:
        print(f"[solver] error: {e}")
        return SolveResponse(success=False, error=str(e))

    if result and result.get("pass"):
        token = result.get("generated_pass_UUID") or result.get("token")
        return SolveResponse(success=True, token=token)
    return SolveResponse(success=False, error="Challenge not solved")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)