captcha_solver/api/routes.py

"""FastAPI routes for the captcha solver API."""

from __future__ import annotations

import time
from threading import Lock

from fastapi import APIRouter, HTTPException

from captcha_solver.config import get_settings
from captcha_solver.models import (
    HealthResponse,
    ModelsResponse,
    SolveRequest,
    SolveResponse,
    StatsResponse,
)
from captcha_solver.solvers.router import SolverRouter
from captcha_solver.utils.cache import SolveCache
from captcha_solver.utils.image import audio_size_hash, perceptual_hash


router = APIRouter()

_state: dict = {
    "router": None,
    "cache": None,
    "stats": {
        "total": 0,
        "by_type": {},
        "by_solver": {},
        "success": 0,
        "elapsed_sum_ms": 0,
        "cache_hits": 0,
        "started_at": time.time(),
    },
    "_stats_lock": Lock(),
}


def get_state() -> dict:
    if _state["router"] is None:
        s = get_settings()
        r = SolverRouter()
        r.init()
        _state["router"] = r
        _state["cache"] = SolveCache(
            ttl_seconds=s.cache_ttl_seconds,
            max_entries=s.cache_max_entries,
        )
    return _state


@router.post("/solve", response_model=SolveResponse)
def solve(req: SolveRequest) -> SolveResponse:
    s = get_settings()
    state = get_state()
    t0 = time.time()
    captcha_type = req.type or "auto"

    if not req.image_base64 and not req.audio_base64 and not req.hint:
        raise HTTPException(status_code=400, detail="must provide image_base64, audio_base64, or hint")

    cache_key = None
    if req.use_cache and s.cache_enabled:
        if req.image_base64:
            try:
                from captcha_solver.utils.image import decode_base64_image
                cache_key = "img:" + perceptual_hash(decode_base64_image(req.image_base64))
            except Exception:
                cache_key = None
        elif req.audio_base64:
            try:
                from captcha_solver.utils.image import decode_base64_audio
                cache_key = "aud:" + audio_size_hash(decode_base64_audio(req.audio_base64))
            except Exception:
                cache_key = None
        if cache_key:
            hit = state["cache"].get(cache_key)
            if hit is not None:
                with state["_stats_lock"]:
                    state["stats"]["total"] += 1
                    state["stats"]["by_type"][captcha_type] = state["stats"]["by_type"].get(captcha_type, 0) + 1
                    state["stats"]["by_solver"][hit.solver] = state["stats"]["by_solver"].get(hit.solver, 0) + 1
                    state["stats"]["success"] += 1
                    state["stats"]["cache_hits"] += 1
                return SolveResponse(
                    success=True,
                    answer=hit.answer,
                    confidence=hit.confidence,
                    solver=hit.solver,
                    elapsed_ms=int((time.time() - t0) * 1000),
                    cache_hit=True,
                )

    router_obj: SolverRouter = state["router"]
    attempt = router_obj.solve(
        captcha_type=captcha_type,
        image_b64=req.image_base64,
        audio_b64=req.audio_base64,
        hint=req.hint,
    )
    elapsed_ms = int((time.time() - t0) * 1000)
    success = bool(attempt.answer) and attempt.confidence > 0

    with state["_stats_lock"]:
        st = state["stats"]
        st["total"] += 1
        st["by_type"][captcha_type] = st["by_type"].get(captcha_type, 0) + 1
        st["by_solver"][attempt.solver_name] = st["by_solver"].get(attempt.solver_name, 0) + 1
        if success:
            st["success"] += 1
        st["elapsed_sum_ms"] += elapsed_ms

    if success and cache_key and s.cache_enabled:
        state["cache"].set(cache_key, attempt.answer, attempt.solver_name, attempt.confidence)

    return SolveResponse(
        success=success,
        answer=attempt.answer or None,
        confidence=attempt.confidence,
        solver=attempt.solver_name,
        elapsed_ms=elapsed_ms,
        cache_hit=False,
        attempts=len(attempt.metadata) + 1,
        error=attempt.error,
    )


@router.get("/health", response_model=HealthResponse)
def health() -> HealthResponse:
    state = get_state()
    engines = state["router"].engine_status() if state["router"] else {}
    status = "ok"
    if any(v == "not_loaded" for v in engines.values()):
        status = "degraded"
    if all(v == "not_loaded" for v in engines.values()):
        status = "down"
    return HealthResponse(
        status=status,
        version="0.1.0",
        engines=engines,
        uptime_s=time.time() - state["stats"]["started_at"],
    )


@router.get("/stats", response_model=StatsResponse)
def stats() -> StatsResponse:
    state = get_state()
    st = state["stats"]
    total = st["total"]
    return StatsResponse(
        total_requests=total,
        by_type=dict(st["by_type"]),
        by_solver=dict(st["by_solver"]),
        success_rate=(st["success"] / total) if total else 0.0,
        avg_elapsed_ms=(st["elapsed_sum_ms"] / total) if total else 0.0,
        cache_hits=st["cache_hits"],
    )


@router.get("/models", response_model=ModelsResponse)
def models() -> ModelsResponse:
    state = get_state()
    engines = state["router"].engine_status() if state["router"] else {}
    loaded = [k for k, v in engines.items() if v == "loaded"]
    available = [
        {"name": "faster-whisper (tiny)", "size_mb": 75, "engine": "whisper", "purpose": "audio captcha"},
        {"name": "faster-whisper (base)", "size_mb": 150, "engine": "whisper", "purpose": "audio captcha (better)"},
        {"name": "Florence-2-base", "size_mb": 1200, "engine": "florence2", "purpose": "OCR + detection"},
        {"name": "Moondream2", "size_mb": 1700, "engine": "moondream2", "purpose": "image grid VQA"},
        {"name": "Qwen2.5-1.5B-Instruct", "size_mb": 1500, "engine": "qwen", "purpose": "math / text reasoning"},
        {"name": "Qwen2-VL 7B (via ollama)", "size_mb": 4500, "engine": "ollama", "purpose": "best vision (optional)"},
    ]
    return ModelsResponse(
        loaded=loaded,
        available=available,
        ollama_enabled=engines.get("ollama") == "loaded",
    )


@router.get("/cache/clear")
def clear_cache() -> dict:
    state = get_state()
    if state["cache"]:
        state["cache"].clear()
    return {"cleared": True}


@router.post("/vqa")
def visual_qa(req: dict) -> dict:
    """Visual Question Answering with Moondream2.

    Request body:
        image_base64: str - base64-encoded image
        question: str - question about the image

    Response:
        answer: str - model's answer
        confidence: float
        solver: str - which model answered
    """
    image_b64 = req.get("image_base64", "")
    question = req.get("question", "")

    if not image_b64:
        raise HTTPException(status_code=400, detail="image_base64 is required")
    if not question:
        raise HTTPException(status_code=400, detail="question is required")

    state = get_state()
    router_obj = state["router"]
    md = router_obj._engines.get("moondream2")
    if not md:
        raise HTTPException(status_code=503, detail="Moondream2 not loaded")
    from captcha_solver.utils.image import decode_base64_image, image_to_pil
    try:
        img = image_to_pil(decode_base64_image(image_b64))
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"bad image: {e}")
    try:
        answer = md.query(img, question)
        return {"answer": answer, "confidence": 0.6, "solver": "moondream2.vqa"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# --- hCaptcha tile classification ---

@router.post("/classify")
def classify_tile(req: dict) -> dict:
    """Classify a single hCaptcha tile against an instruction.

    Request body:
        image_base64: str - base64-encoded tile image
        instruction: str - hCaptcha instruction (e.g. "Find all items made by people")

    Response:
        match: bool - whether the tile matches the instruction
        confidence: float - classification confidence
        caption: str - Florence-2 caption of the tile
        solver: str - which solver was used
    """
    from captcha_solver.solvers.hcaptcha_solver import classify_tile as _classify

    image_b64 = req.get("image_base64", "")
    instruction = req.get("instruction", "")

    if not image_b64:
        raise HTTPException(status_code=400, detail="image_base64 is required")
    if not instruction:
        raise HTTPException(status_code=400, detail="instruction is required")

    state = get_state()
    router_obj = state["router"]
    # Build SolveContext from router's engines
    from captcha_solver.solvers.base import SolveContext
    ctx = SolveContext(
        whisper=router_obj._engines.get("whisper"),
        florence=router_obj._engines.get("florence2"),
        moondream=router_obj._engines.get("moondream2"),
        qwen=router_obj._engines.get("qwen"),
        ollama=router_obj._engines.get("ollama"),
    )
    return _classify(image_b64, instruction, ctx)


# --- Learning stats ---
try:
    from captcha_solver.learning.db import LearningDB
    _learning_db = LearningDB()
except ImportError:
    _learning_db = None


@router.get("/stats")
def get_learning_stats() -> dict:
    if not _learning_db:
        return {"enabled": False, "error": "learning module not available"}
    summary = _learning_db.summary()
    rankings = _learning_db.get_solver_ranking()
    return {
        "enabled": True,
        "summary": summary,
        "rankings": rankings,
        "recent_failures": _learning_db.get_recent_failures(5),
    }