"""Shared Vision AI pipeline — screenshot, analyze, cache, retry, beep.

Used by AnswerDisplay (Ctrl+F9), AnswerWindow (Ctrl+F3), MouseClicker (Ctrl+Shift+M).
Eliminates code duplication across all three answer functions.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import re
import sys
import threading
import time
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ..config import Config

log = logging.getLogger(__name__)

# JSON object extractor — greedily finds last JSON object in response
JSON_RE = re.compile(r"\{[^{}]*\}", re.DOTALL)

# Simple answer prompt — optimized for minicpm-v (local Ollama model)
# Strategy: ask for letter FIRST to avoid truncation cutting off the answer.
ANSWER_PROMPT = """\
You are an expert exam solver. Look at this exam question screenshot carefully.

Step 1: Read the question and all answer choices (A, B, C, D).
Step 2: Think which answer is correct.
Step 3: Reply with ONLY the correct letter: A, B, C, or D.

Important: Output ONLY one letter. No explanation. No punctuation. Just the letter.
"""

# OCR Reader prompt — extract all visible text from screenshot
OCR_READER_PROMPT = """\
Read and extract ALL text visible in this image. Return the complete text exactly as shown, preserving the original language (Chinese/English). Do not translate, do not answer questions, just extract the text.
Include every word, sentence, paragraph, heading, label, and option you can see.
"""

# Mouse clicker prompt (answer + position)
MOUSE_PROMPT = """\
You are an expert exam solver. Look at this screenshot of a test/exam.

1. Find the multiple-choice question and read ALL answer options carefully
2. Solve it — verify your answer is correct
3. Look at WHERE the correct answer option is located on screen:
   - y_pct: how far down from the TOP of the image (0.0 = very top, 1.0 = very bottom)
   - x_pct: how far from the LEFT edge (0.0 = left edge, 1.0 = right edge)
   Estimate the CENTER of the correct answer option text/radio button.

Reply with ONLY one line of JSON, nothing else:
{{"found":true,"answer":"B","y_pct":0.45,"x_pct":0.3}}

Example positions:
- First option near top of page: y_pct ~0.3-0.4
- Second option: y_pct ~0.4-0.5
- Third option: y_pct ~0.5-0.6
- Fourth option near bottom: y_pct ~0.6-0.7
- Options usually start from left: x_pct ~0.15-0.35

If no question visible: {{"found":false,"answer":"","y_pct":0,"x_pct":0}}
Do NOT add any text before or after the JSON.
"""

# Verification prompt template (two-pass AI verification)
VERIFY_PROMPT_TEMPLATE = (
    "你是考试审核专家。仔细看这张截图。上一位专家认为正确答案是{answer}。"
    "请你独立分析题目和所有选项，验证这个答案是否正确。"
    '只回复JSON：{{"found":true,"answer":"A","verified":true}}'
)


# ===================================================================== #
# Answer cache — maps screenshot hash → (answer_dict, timestamp)
# ===================================================================== #

_answer_cache: dict[str, tuple[dict, float]] = {}
_answer_cache_lock = threading.Lock()
_ANSWER_CACHE_TTL = 10.0   # seconds — short to avoid stale answers on new questions
_ANSWER_CACHE_MAX = 20      # max entries

# ===================================================================== #
# Text-based question cache (improvement 1) — maps question text hash → answer dict
# ===================================================================== #

_question_text_cache: dict[str, dict] = {}
_question_text_cache_lock = threading.Lock()
_QUESTION_TEXT_CACHE_MAX = 50

# ===================================================================== #
# OCR image cache — maps image hash → (ocr_text, timestamp)
# Skips Stage 1 vision model call entirely on repeated screenshots.
# ===================================================================== #

_ocr_image_cache: dict[str, tuple[str, float]] = {}
_ocr_image_cache_lock = threading.Lock()
_OCR_IMAGE_CACHE_TTL = 15.0   # 15 seconds — short to avoid stale OCR on new questions
_OCR_IMAGE_CACHE_MAX = 30


def _ocr_image_cache_get(img_b64: str) -> str | None:
    """Return cached OCR text if same screenshot seen before (within TTL)."""
    h = _get_screenshot_hash(img_b64)
    with _ocr_image_cache_lock:
        entry = _ocr_image_cache.get(h)
        if entry is not None:
            text, ts = entry
            if (time.monotonic() - ts) < _OCR_IMAGE_CACHE_TTL:
                log.info("[two-stage] OCR image cache HIT — skipping Stage 1 (hash=%s…)", h[:12])
                return text
            del _ocr_image_cache[h]
    return None


def _ocr_image_cache_put(img_b64: str, text: str) -> None:
    """Store OCR result in image cache, evicting oldest if at capacity."""
    h = _get_screenshot_hash(img_b64)
    now = time.monotonic()
    with _ocr_image_cache_lock:
        while len(_ocr_image_cache) >= _OCR_IMAGE_CACHE_MAX:
            oldest = min(_ocr_image_cache, key=lambda k: _ocr_image_cache[k][1])
            del _ocr_image_cache[oldest]
        _ocr_image_cache[h] = (text, now)
        log.info("[two-stage] OCR image cache PUT (hash=%s…, size=%d)", h[:12], len(_ocr_image_cache))


def _get_question_text_hash(text: str) -> str:
    """Compute SHA-256 hash of extracted question text (normalized)."""
    normalized = " ".join(text.split()).lower()
    return hashlib.sha256(normalized.encode("utf-8")).hexdigest()


_QUESTION_TEXT_CACHE_TTL = 30.0  # 30 seconds — avoid stale answers


def _question_text_cache_get(text: str) -> dict | None:
    """Return cached answer dict if same question text seen before (within TTL), else None."""
    h = _get_question_text_hash(text)
    with _question_text_cache_lock:
        entry = _question_text_cache.get(h)
        if entry is not None:
            result, ts = entry
            if (time.monotonic() - ts) < _QUESTION_TEXT_CACHE_TTL:
                log.info("[pipeline] Question text cache HIT (hash=%s…)", h[:12])
                return result
            del _question_text_cache[h]
    return None


def _question_text_cache_put(text: str, result: dict) -> None:
    """Store answer in text cache with TTL, evicting oldest if over max size."""
    h = _get_question_text_hash(text)
    now = time.monotonic()
    with _question_text_cache_lock:
        if len(_question_text_cache) >= _QUESTION_TEXT_CACHE_MAX:
            oldest_key = next(iter(_question_text_cache))
            del _question_text_cache[oldest_key]
        _question_text_cache[h] = (result, now)
        log.info("[pipeline] Question text cache PUT (hash=%s…, size=%d)",
                 h[:12], len(_question_text_cache))


def _get_screenshot_hash(b64: str) -> str:
    """Compute a reliable hash of a base64 screenshot string.

    Uses full content to avoid hash collisions between different screenshots.
    """
    return hashlib.sha256(b64.encode("utf-8")).hexdigest()


def _answer_cache_get(b64: str) -> dict | None:
    """Return cached answer dict if cache hit within TTL, else None."""
    h = _get_screenshot_hash(b64)
    with _answer_cache_lock:
        entry = _answer_cache.get(h)
        if entry is not None:
            result, ts = entry
            if (time.monotonic() - ts) < _ANSWER_CACHE_TTL:
                log.info("[pipeline] Answer cache HIT (hash=%s…)", h[:12])
                return result
            else:
                # Expired — remove
                del _answer_cache[h]
    return None


def _answer_cache_put(b64: str, result: dict) -> None:
    """Store answer in cache, evicting oldest if over max size."""
    h = _get_screenshot_hash(b64)
    now = time.monotonic()
    with _answer_cache_lock:
        # Evict oldest entries if at capacity
        while len(_answer_cache) >= _ANSWER_CACHE_MAX:
            oldest_key = min(_answer_cache, key=lambda k: _answer_cache[k][1])
            del _answer_cache[oldest_key]
        _answer_cache[h] = (result, now)
        log.info("[pipeline] Answer cache PUT (hash=%s…, size=%d)",
                 h[:12], len(_answer_cache))


# ===================================================================== #
# Screenshot cache — shared between F3/F9/Shift+M
# ===================================================================== #

class _ScreenshotCache:
    """Thread-safe screenshot cache with TTL."""

    def __init__(self, ttl: float = 3.0) -> None:
        self._lock = threading.Lock()
        self._b64: str = ""
        self._orig_w: int = 0
        self._orig_h: int = 0
        self._res_w: int = 0
        self._res_h: int = 0
        self._ts: float = 0.0
        self._ttl = ttl

    def get(self) -> tuple[str, int, int, int, int] | None:
        """Return cached (b64, orig_w, orig_h, res_w, res_h) or None if expired."""
        with self._lock:
            if self._b64 and (time.monotonic() - self._ts) < self._ttl:
                log.info("[pipeline] Using cached screenshot (age=%.1fs)",
                         time.monotonic() - self._ts)
                return self._b64, self._orig_w, self._orig_h, self._res_w, self._res_h
            return None

    def put(self, b64: str, orig_w: int, orig_h: int, res_w: int, res_h: int) -> None:
        """Store screenshot in cache."""
        with self._lock:
            self._b64 = b64
            self._orig_w = orig_w
            self._orig_h = orig_h
            self._res_w = res_w
            self._res_h = res_h
            self._ts = time.monotonic()

    def clear(self) -> None:
        """Invalidate cache."""
        with self._lock:
            self._b64 = ""
            self._ts = 0.0


# Global cache instance (shared between all answer functions)
screenshot_cache = _ScreenshotCache(ttl=3.0)


# ===================================================================== #
# DPI awareness
# ===================================================================== #

_dpi_done = False


def ensure_dpi_aware() -> None:
    """Make process DPI-aware on Windows (idempotent)."""
    global _dpi_done
    if _dpi_done or sys.platform != "win32":
        return
    try:
        import ctypes
        ctypes.windll.shcore.SetProcessDpiAwareness(2)  # type: ignore[attr-defined]
        _dpi_done = True
    except Exception:
        pass


# ===================================================================== #
# Screenshot capture
# ===================================================================== #

def capture_screenshot(config: Config, config_prefix: str) -> tuple[str, int, int, int, int]:
    """Take screenshot, resize for AI. Uses cache if available.

    Args:
        config: App config.
        config_prefix: Config key prefix (e.g. "answer_display", "mouse_clicker").

    Returns:
        (base64, orig_w, orig_h, resized_w, resized_h) or ("", 0, 0, 0, 0) on error.
    """
    ensure_dpi_aware()

    # Check cache first
    cached = screenshot_cache.get()
    if cached:
        return cached

    try:
        from ..screenshot import capture_screen
        from ..utils.image_utils import image_to_base64, resize_image

        img = capture_screen()
        orig_w, orig_h = img.size

        # Cap at 1920x1080; smaller screenshots pass through unchanged.
        img_resized = resize_image(img)
        res_w, res_h = img_resized.size

        b64 = image_to_base64(img_resized)
        log.info("[pipeline] Screenshot %dx%d resized=%dx%d",
                 orig_w, orig_h, res_w, res_h)

        # Store in cache
        screenshot_cache.put(b64, orig_w, orig_h, res_w, res_h)

        return b64, orig_w, orig_h, res_w, res_h

    except Exception as exc:
        log.error("[pipeline] Screenshot error: %s", exc)
        return "", 0, 0, 0, 0


# ===================================================================== #
# Vision AI analysis with retry
# ===================================================================== #

def analyze_screenshot(
    config: Config,
    img_b64: str,
    prompt: str,
    *,
    config_prefix: str = "answer_display",
    max_tokens: int = 64,
    retries: int = 2,
    required_key: str = "answer",
) -> dict | None:
    """Send screenshot to Vision AI. Returns parsed JSON dict or None.

    Includes retry with exponential backoff and provider fallback chain:
    if the current provider fails with connection errors, tries the next
    available provider (proxy → aliyun → 88api).

    Args:
        config: App config.
        img_b64: Base64-encoded screenshot.
        prompt: AI prompt.
        config_prefix: Config key prefix for timeout.
        max_tokens: Max tokens for AI response.
        retries: Number of retry attempts on failure.
        required_key: Required key in JSON response ("answer" or "found").

    Returns:
        Parsed JSON dict or None.
    """
    if not img_b64:
        log.error("[pipeline] analyze_screenshot called with empty img_b64 — aborting")
        return None

    from ..ai.vision_adapter import VisionAdapter

    timeout = float(config.get(f"{config_prefix}.ai_timeout", 30.0))

    # Vision fallback chain: local first, then online (through proxy).
    adapter = VisionAdapter(config, timeout=timeout)

    # Build fallback chain: online first (fast 2-5s) → ollama → llamacpp (slow 40-90s)
    fallback_chain: list[str] = []

    # Online vision providers FIRST (fast, routed through proxy + keepalive)
    # These work even under CSCA thanks to NetworkProtector
    fallback_chain.append(config.vision_provider)  # configured provider first
    online_vision = [p for p in ["proxy", "aliyun", "88api"] if p != config.vision_provider]
    fallback_chain.extend(online_vision)

    # Ollama as middle option (localhost)
    fallback_chain.append("ollama")
    try:
        from ..seb.ollama_guard import ensure_ollama_running
        ensure_ollama_running()
    except Exception:
        pass

    # llamacpp as LAST resort (in-process but very slow on CPU: 40-90s per image)
    try:
        from ..ai.provider_llamacpp import is_llamacpp_available, are_models_present
        if is_llamacpp_available() and are_models_present(vision=True, text=False):
            fallback_chain.append("llamacpp")
    except ImportError:
        pass

    # Deduplicate while preserving order
    seen: set[str] = set()
    unique_chain: list[str] = []
    for p in fallback_chain:
        if p not in seen:
            seen.add(p)
            unique_chain.append(p)
    fallback_chain = unique_chain
    log.info("[pipeline] Vision fallback chain: %s", " → ".join(fallback_chain))

    for provider_idx, provider_name in enumerate(fallback_chain):
        # Temporarily override vision provider for this attempt
        original_provider = config.vision_provider
        if provider_idx > 0:
            log.info("[pipeline] Falling back to provider: %s (Ollama likely blocked by CSCA)", provider_name)
            config._data.setdefault("vision", {})["provider"] = provider_name

        try:
            result = _try_provider(
                config, img_b64, prompt,
                timeout=timeout,
                max_tokens=max_tokens,
                retries=retries if provider_idx == 0 else 1,
                required_key=required_key,
            )
            if result is not None:
                return result
        finally:
            # Restore original provider
            config._data.setdefault("vision", {})["provider"] = original_provider

    log.error("[Vision] AI недоступен. Запустите setup_llamacpp.bat для установки локального AI.")
    return None


# ===================================================================== #
# Two-pass AI verification
# ===================================================================== #

def analyze_screenshot_verified(
    config: Config,
    img_b64: str,
    prompt: str,
    *,
    config_prefix: str = "answer_display",
    max_tokens: int = 64,
    retries: int = 2,
    required_key: str = "answer",
    verify: bool = True,
) -> dict | None:
    """Analyze screenshot with optional two-pass verification and answer caching.

    1. Check answer cache — return immediately on hit.
    2. Call analyze_screenshot() to get the first answer.
    3. If verify=True and answer found, call analyze_screenshot() again with
       a verification prompt. If both answers match, return with high confidence.
       If they differ, return the second (verification) answer.
    4. Cache the final result.

    Args:
        config: App config.
        img_b64: Base64-encoded screenshot.
        prompt: AI prompt for the first pass.
        config_prefix: Config key prefix for timeout.
        max_tokens: Max tokens for AI response.
        retries: Number of retry attempts on failure.
        required_key: Required key in JSON response.
        verify: If True, run a second verification pass.

    Returns:
        Parsed JSON dict or None.
    """
    if not img_b64:
        log.error("[pipeline] analyze_screenshot_verified called with empty img_b64")
        return None

    # --- Answer cache check ---
    cached = _answer_cache_get(img_b64)
    if cached is not None:
        return cached

    # --- First pass ---
    result = analyze_screenshot(
        config, img_b64, prompt,
        config_prefix=config_prefix,
        max_tokens=max_tokens,
        retries=retries,
        required_key=required_key,
    )
    if result is None:
        return None

    first_answer = extract_answer(result)
    if not first_answer or not verify:
        # No answer found or verification disabled — cache and return as-is
        if first_answer:
            _answer_cache_put(img_b64, result)
        return result

    # --- Second pass (verification) ---
    verify_prompt = VERIFY_PROMPT_TEMPLATE.format(answer=first_answer)
    log.info("[pipeline] Verification pass for answer=%s", first_answer)

    verify_result = analyze_screenshot(
        config, img_b64, verify_prompt,
        config_prefix=config_prefix,
        max_tokens=max_tokens,
        retries=1,           # fewer retries for verification
        required_key=required_key,
    )

    if verify_result is None:
        # Verification call failed — trust first answer
        log.warning("[pipeline] Verification failed — using first answer: %s", first_answer)
        _answer_cache_put(img_b64, result)
        return result

    second_answer = extract_answer(verify_result)

    if first_answer == second_answer:
        # Both agree — high confidence
        log.info("[pipeline] Verification MATCH: %s (high confidence)", first_answer)
        result["verified"] = True
        _answer_cache_put(img_b64, result)
        return result
    else:
        # Disagreement — trust the verification (second) pass
        log.info("[pipeline] Verification DIFFERS: first=%s, second=%s — using second",
                 first_answer, second_answer)
        verify_result["verified"] = True
        _answer_cache_put(img_b64, verify_result)
        return verify_result


# ===================================================================== #
# Two-stage strategy: Vision OCR → Text Solve
# ===================================================================== #

# OCR prompt — extract ALL text from screenshot including passage/article
OCR_EXTRACT_PROMPT = """\
Read this exam screenshot carefully. Extract ALL visible text from top to bottom.

IMPORTANT: Include the FULL reading passage or article text (if any) that appears ABOVE the question. This context is required to answer the question correctly.

Format your output as:
[Passage/Article/Context text here — copy EVERY sentence verbatim, do NOT skip any]

Question: [the question text]

If it is a multiple-choice question (A/B/C/D), format:
A) [option A text]
B) [option B text]
C) [option C text]
D) [option D text]
QUESTION_TYPE: multiple_choice

If it is a True/False question, format:
True) [True option text if shown]
False) [False option text if shown]
QUESTION_TYPE: true_false

Otherwise:
QUESTION_TYPE: other

If no question visible, reply: NO QUESTION
"""

# Few-shot examples for Stage 2 solve prompt — teaches model the expected format
_FEW_SHOT_MULTIPLE_CHOICE = """\
Example 1:
Passage: The Ming dynasty emperors built the Great Wall to defend against northern invasions.
Question: What was the main purpose of the Great Wall?
A) To mark the empire boundary  B) To defend against northern invasions  C) To collect taxes  D) To connect trade routes
Answer: B

Example 2:
Question: 下列哪个词语是动词？
A) 美丽  B) 书本  C) 跑步  D) 高兴
Answer: C

Now solve the following exam question. Reply with ONLY one letter (A, B, C, or D):
"""

_FEW_SHOT_TRUE_FALSE = """\
Example 1:
Statement: The capital of China is Shanghai.
Answer: False

Example 2:
Statement: Water freezes at 0 degrees Celsius.
Answer: True

Now answer True or False for the following. Reply with ONLY True or False:
"""


def _detect_question_type(ocr_text: str) -> str:
    """Detect question type from OCR-extracted text.

    Returns: 'multiple_choice', 'true_false', or 'other'.
    """
    upper = ocr_text.upper()
    if "QUESTION_TYPE: MULTIPLE_CHOICE" in upper:
        return "multiple_choice"
    if "QUESTION_TYPE: TRUE_FALSE" in upper:
        return "true_false"
    # Heuristic fallback: check for A) B) C) D) pattern
    if re.search(r'\bA\)', ocr_text) and re.search(r'\bB\)', ocr_text):
        return "multiple_choice"
    # Check for true/false keywords
    if re.search(r'\b(true|false|对|错|正确|错误|是|否)\b', ocr_text, re.IGNORECASE):
        return "true_false"
    return "other"


def _append_answer_history(question_text: str, answer: str, confidence: str) -> None:
    """Append a line to ~/.muravey-china/answer_history.log (thread-safe, never crashes)."""
    import datetime
    try:
        log_dir = os.path.expanduser("~/.muravey-china")
        os.makedirs(log_dir, exist_ok=True)
        log_path = os.path.join(log_dir, "answer_history.log")
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        q_snippet = question_text[:100].replace("\n", " ")
        line = f"{timestamp} | Q: {q_snippet} | A: {answer} | conf: {confidence}\n"
        with _answer_history_lock:
            with open(log_path, "a", encoding="utf-8") as f:
                f.write(line)
    except Exception as exc:
        log.debug("[two-stage] Failed to write answer history: %s", exc)


_answer_history_lock = threading.Lock()


def analyze_two_stage(
    config: "Config",
    img_b64: str,
    *,
    config_prefix: str = "answer_window",
    max_tokens: int = 1024,
) -> dict | None:
    """Two-stage answer pipeline: Vision OCR → Text Solve.

    Stage 1: Vision model reads screenshot, extracts question + options as text.
    Stage 2: Text model (qwen2.5:7b) solves the question from extracted text.

    This is more accurate than asking the vision model to solve directly,
    because text models are better at reasoning.

    Returns:
        dict with {"found": True, "answer": "A", "confidence": "high"} or None on failure.
    """
    if not img_b64:
        return None

    # Stage 1: Vision OCR — extract question text
    log.info("[two-stage] Stage 1: extracting question text via llamacpp (in-process)")
    ocr_result_raw = ""

    from ..ai.vision_adapter import VisionAdapter
    timeout = float(config.get(f"{config_prefix}.ai_timeout", 30.0))

    # Check OCR image cache first — skip vision model for same screenshot
    _cached_ocr = _ocr_image_cache_get(img_b64)
    if _cached_ocr is not None:
        ocr_result_raw = _cached_ocr
        log.info("[two-stage] Using OCR image cache — Stage 1 skipped")
    else:
        # Skip preprocessing on CPU to save time — use raw image directly
        _preprocessed_b64 = img_b64

        # PRIMARY: online vision providers (fast 2-5s, through proxy)
        ocr_fallback = [config.vision_provider]
        for p in ["proxy", "aliyun", "88api"]:
            if p not in ocr_fallback:
                ocr_fallback.append(p)
        ocr_fallback.append("ollama")

        for fb_provider in ocr_fallback:
            if ocr_result_raw.strip():
                break

            if fb_provider == "ollama":
                try:
                    from ..seb.ollama_guard import ensure_ollama_running
                    ensure_ollama_running()
                except Exception:
                    pass

            adapter = VisionAdapter(config, timeout=timeout)
            cancel = threading.Event()

            try:
                original_provider = config.vision_provider
                config._data.setdefault("vision", {})["provider"] = fb_provider
                try:
                    model = config.vision_model if hasattr(config, "vision_model") else None
                    if fb_provider == "ollama" and not model:
                        model = "qwen2.5-vl:7b"
                    kwargs = dict(
                        cancel_event=cancel,
                        stream=False,
                        system_prompt="You are an OCR assistant. Extract ALL visible text from the image exactly as written. Include every word, sentence, and paragraph you see. Do NOT summarize or skip anything.",
                        temperature=0.0,
                        max_tokens=512,
                    )
                    if model:
                        kwargs["model"] = model
                    for chunk in adapter.analyze(_preprocessed_b64, OCR_EXTRACT_PROMPT, **kwargs):
                        ocr_result_raw += chunk
                    if ocr_result_raw.strip():
                        log.info("[two-stage] Stage 1 OCR via %s: %d chars", fb_provider, len(ocr_result_raw))
                finally:
                    config._data.setdefault("vision", {})["provider"] = original_provider
            except Exception as exc:
                log.warning("[two-stage] Stage 1 %s failed: %s", fb_provider, exc)

        # LAST RESORT: llamacpp (slow 40-90s but works fully offline)
        if not ocr_result_raw.strip():
            try:
                from ..ai.provider_llamacpp import get_llamacpp_provider
                lcpp = get_llamacpp_provider()
                if lcpp is not None and lcpp.vision_ready:
                    log.info("[two-stage] Stage 1 LAST RESORT: using llama-cpp-python (slow, no network)")
                    ocr_result_raw = lcpp.generate(
                        OCR_EXTRACT_PROMPT,
                        image_base64=_preprocessed_b64,
                        system_prompt="You are an OCR assistant. Extract ALL visible text from the image exactly as written.",
                        max_tokens=512,
                        temperature=0.0,
                    )
            except Exception as lcpp_exc:
                log.warning("[two-stage] Stage 1 llamacpp failed: %s", lcpp_exc)

        if not ocr_result_raw.strip():
            log.error("[two-stage] Stage 1 failed with all providers")
            return None

        # Cache OCR result for future identical screenshots
        if ocr_result_raw.strip():
            _ocr_image_cache_put(img_b64, ocr_result_raw)

    log.info("[two-stage] OCR extracted: %r", ocr_result_raw[:300])

    if not ocr_result_raw.strip() or "NO QUESTION" in ocr_result_raw.upper():
        log.info("[two-stage] No question found in screenshot")
        return None

    # Improvement 1: Text-based question cache check (after OCR, before Stage 2)
    cached = _question_text_cache_get(ocr_result_raw)
    if cached is not None:
        log.info("[two-stage] Question text cache HIT — skipping Stage 2")
        return cached

    # Improvement 2: Detect question type from OCR text
    question_type = _detect_question_type(ocr_result_raw)
    log.info("[two-stage] Question type detected: %s", question_type)

    # Stage 2: Solve via fast text model (small = stays in RAM alongside vision model)
    # Priority: config override → deepseek-chat (fast online)
    text_model = config.get("answer_window.text_model", "deepseek-chat")
    log.info("[two-stage] Stage 2: solving via text model (%s)", text_model)

    # Build solve prompt based on question type — with few-shot examples for accuracy
    if question_type == "true_false":
        solve_prompt = (
            _FEW_SHOT_TRUE_FALSE
            + "\n"
            + ocr_result_raw.strip()
        )
        system_prompt_stage2 = "You are an exam answer machine. Reply with ONLY 'True' or 'False'. Nothing else."
    elif question_type == "multiple_choice":
        solve_prompt = (
            _FEW_SHOT_MULTIPLE_CHOICE
            + "\n"
            + ocr_result_raw.strip()
        )
        system_prompt_stage2 = "You are an exam answer machine. Reply with ONLY one letter: A, B, C, or D. Nothing else."
    else:
        solve_prompt = (
            _FEW_SHOT_MULTIPLE_CHOICE
            + "\nIf the question is True/False, reply True or False instead.\n\n"
            + ocr_result_raw.strip()
        )
        system_prompt_stage2 = "You are an exam answer machine. Reply with ONLY the answer (one letter A/B/C/D, or True/False). Nothing else."

    def _call_text_model() -> tuple[str | None, str]:
        """Call text model, return (answer, confidence) or (None, '').

        Uses ChinaAdapter (online providers: deepseek/aliyun first, then local fallbacks).
        """
        text_answer_raw: str | None = None

        try:
            from ..ai.china_adapter import ChinaAdapter
            from ..ai.base import AIRequest

            adapter = ChinaAdapter(config)
            request = AIRequest(
                prompt=solve_prompt,
                system_prompt=system_prompt_stage2,
                temperature=0.0,
                max_tokens=8,
                model=text_model,
            )

            response = adapter.ask(request)
            if response and response.text:
                text_answer_raw = response.text.strip()
        except Exception as exc:
            log.error("[two-stage] ChinaAdapter text model failed: %s", exc)

        try:
            if text_answer_raw:
                text_answer = text_answer_raw.strip().upper()
                log.info("[two-stage] Answer from %s: %r", text_model, text_answer)

                # True/False handling
                if question_type == "true_false":
                    if "TRUE" in text_answer or "对" in text_answer or "正确" in text_answer:
                        return "True", "high"
                    if "FALSE" in text_answer or "错" in text_answer:
                        return "False", "high"
                    # Fallback for true/false
                    if text_answer.startswith("T"):
                        return "True", "medium"
                    if text_answer.startswith("F"):
                        return "False", "medium"
                    return None, ""

                # Multiple choice handling
                clean = text_answer.rstrip('.').strip()
                if len(clean) == 1 and clean in "ABCD":
                    return clean, "high"

                # Try patterns (medium confidence)
                for pat in [
                    r'(?:答案|正确答案|选|answer|correct)\s*[:：.]?\s*([A-D])',
                    r'^([A-D])(?:\s|$|\.)',
                ]:
                    m = re.search(pat, text_answer, re.IGNORECASE)
                    if m:
                        return m.group(1).upper(), "medium"

                # First A-D letter (medium confidence)
                m = re.search(r'(?<![a-zA-Z])([A-D])(?![a-zA-Z])', text_answer)
                if m:
                    return m.group(1).upper(), "medium"

                # Last resort: last standalone letter (low confidence)
                all_letters = re.findall(r'(?<![a-zA-Z])([A-D])(?![a-zA-Z])', text_answer)
                if all_letters:
                    return all_letters[-1], "low"
        except Exception as exc:
            log.error("[two-stage] Stage 2 call failed: %s", exc)
        return None, ""

    # 1x call for speed on CPU — skip majority voting to save ~10s
    votes: list[tuple[str, str]] = []
    for _vote_i in range(1):
        _ans, _conf = _call_text_model()
        if _ans is not None:
            votes.append((_ans, _conf))
            log.info("[two-stage] Vote %d/%d: %s (conf=%s)", _vote_i + 1, 3, _ans, _conf)

    if not votes:
        return None

    # Tally votes
    from collections import Counter as _Counter
    _counts = _Counter(v[0] for v in votes)
    _best_answer, _best_count = _counts.most_common(1)[0]

    if _best_count == 3:
        # All 3 agree — maximum confidence
        final_answer = _best_answer
        final_confidence = "verified"
        log.info("[two-stage] 3/3 unanimous: %s — verified", final_answer)
    elif _best_count == 2:
        # 2 out of 3 agree — high confidence
        final_answer = _best_answer
        final_confidence = "high"
        log.info("[two-stage] 2/3 majority: %s — high confidence", final_answer)
    else:
        # All different — use answer with best individual confidence
        _conf_rank = {"high": 3, "medium": 2, "verified": 4, "low": 1}
        votes_sorted = sorted(votes, key=lambda v: _conf_rank.get(v[1], 0), reverse=True)
        final_answer, final_confidence = votes_sorted[0]
        log.info("[two-stage] All differ — using best-confidence: %s (%s)", final_answer, final_confidence)

    result = {
        "found": True,
        "answer": final_answer,
        "two_stage": True,
        "confidence": final_confidence,  # improvement 5
    }

    # Store in text-based cache (improvement 1)
    _question_text_cache_put(ocr_result_raw, result)

    # Improvement 6: Append to answer history log
    threading.Thread(
        target=_append_answer_history,
        args=(ocr_result_raw, final_answer, final_confidence),
        daemon=True,
    ).start()

    return result


def _try_provider(
    config: Config,
    img_b64: str,
    prompt: str,
    *,
    timeout: float,
    max_tokens: int,
    retries: int,
    required_key: str,
) -> dict | None:
    """Try a single vision provider with retries. Returns parsed dict or None."""
    from ..ai.vision_adapter import VisionAdapter

    for attempt in range(1, retries + 2):
        cancel = threading.Event()
        try:
            adapter = VisionAdapter(config, timeout=timeout)

            raw = ""
            # Always use non-streaming for answer detection — faster response.
            use_stream = False
            for chunk in adapter.analyze(
                img_b64, prompt,
                cancel_event=cancel,
                stream=use_stream,
                model=config.vision_model if hasattr(config, "vision_model") else "qwen2.5-vl:7b",
                temperature=0.0,
                max_tokens=max_tokens,
            ):
                raw += chunk

            log.info("[pipeline] AI raw (attempt %d): %r", attempt, raw[:300])

            # Connection error in response text — don't parse, retry
            if "[Vision error:" in raw and "Connection error" in raw:
                log.warning("[pipeline] Connection error in response (attempt %d)", attempt)
                if attempt <= retries:
                    delay = min(1.0 * (2.0 ** (attempt - 1)), 10.0)
                    log.info("[pipeline] Retrying in %.1fs...", delay)
                    time.sleep(delay)
                continue

            # Find JSON with required key
            matches = JSON_RE.findall(raw)
            for candidate in reversed(matches):
                try:
                    parsed = json.loads(candidate)
                    if required_key in parsed:
                        return parsed
                except json.JSONDecodeError:
                    continue

            # Fallback: extract answer letter from free text (minicpm-v often skips JSON)
            if required_key == "answer" and raw.strip():
                # Strategy: if the model replied with JUST a letter (or letter + period), use it
                clean = raw.strip().rstrip('.').strip()
                if len(clean) == 1 and clean in "ABCD":
                    log.info("[pipeline] Fallback single letter: '%s' (attempt %d)", clean, attempt)
                    return {"found": True, "answer": clean}

                # Look for conclusion patterns (Chinese/English) — these indicate the AI's ANSWER
                conclusion_patterns = [
                    r'(?:答案|正确答案|选|应选|应该选|所以选|因此选|故选)\s*[:：]?\s*([A-D])',
                    r'(?:answer|correct answer|choose|select|therefore)\s*[:.]?\s*([A-D])',
                    r'([A-D])\s*(?:是正确|为正确|is correct|is the correct|is the answer)',
                    r'(?:正确|correct)[^A-D]{0,20}([A-D])',
                ]
                for pat in conclusion_patterns:
                    m = re.search(pat, raw, re.IGNORECASE)
                    if m:
                        letter = m.group(1).upper()
                        log.info("[pipeline] Fallback conclusion: extracted '%s' (attempt %d)", letter, attempt)
                        return {"found": True, "answer": letter}

                # If response starts with a letter (model followed "just one letter" instruction partially)
                first_char = raw.strip()[0].upper() if raw.strip() else ""
                if first_char in "ABCD":
                    log.info("[pipeline] Fallback first char: '%s' (attempt %d)", first_char, attempt)
                    return {"found": True, "answer": first_char}

                # Last resort: find the LAST standalone A-D letter in text
                # (the model's conclusion is usually at the end)
                all_letters = re.findall(r'(?<![a-zA-Z])([A-D])(?![a-zA-Z])', raw)
                if all_letters:
                    letter = all_letters[-1]
                    log.info("[pipeline] Fallback last letter: extracted '%s' (attempt %d)", letter, attempt)
                    return {"found": True, "answer": letter}

            log.warning("[pipeline] No JSON with '%s' in response (attempt %d): %r",
                        required_key, attempt, raw[:200])

        except Exception as exc:
            cancel.set()  # signal adapter generator to stop streaming
            log.error("[pipeline] AI error (attempt %d): %s", attempt, exc)

        # Retry with exponential backoff
        if attempt <= retries:
            delay = min(1.0 * (2.0 ** (attempt - 1)), 10.0)
            log.info("[pipeline] Retrying AI request in %.1fs (attempt %d/%d)...",
                     delay, attempt + 1, retries + 1)
            time.sleep(delay)

    return None


# ===================================================================== #
# Parse answer from AI result
# ===================================================================== #

def extract_answer(result: dict) -> str:
    """Extract and normalize answer letter from AI result.

    Returns uppercase single letter (A-D) or empty string.
    """
    answer = str(result.get("answer") or "").strip().upper()
    if not answer:
        return ""
    # Only keep first letter
    return answer[0]


def ocr_read_screenshot(
    config: "Config",
    img_b64: str,
    *,
    config_prefix: str = "answer_window",
    max_tokens: int = 2048,
) -> str | None:
    """Read all text from a screenshot using the vision model.

    Returns:
        Extracted text string, or None on failure.
    """
    if not img_b64:
        return None

    from ..ai.vision_adapter import VisionAdapter
    timeout = float(config.get(f"{config_prefix}.ai_timeout", 30.0))

    raw = ""

    # PRIMARY: online vision providers (fast 2-5s, through proxy)
    ocr_fallback = [config.vision_provider]
    for p in ["proxy", "aliyun", "88api"]:
        if p not in ocr_fallback:
            ocr_fallback.append(p)
    ocr_fallback.append("ollama")

    if not raw.strip():

        for fb_provider in ocr_fallback:
            if raw.strip():
                break

            if fb_provider == "ollama":
                try:
                    from ..seb.ollama_guard import ensure_ollama_running
                    ensure_ollama_running()
                except Exception:
                    pass

            adapter = VisionAdapter(config, timeout=timeout)
            cancel = threading.Event()

            try:
                original_provider = config.vision_provider
                config._data.setdefault("vision", {})["provider"] = fb_provider
                try:
                    model = config.vision_model if hasattr(config, "vision_model") else None
                    if fb_provider == "ollama" and not model:
                        model = "qwen2.5-vl:7b"
                    kwargs = dict(
                        cancel_event=cancel, stream=False,
                        temperature=0.0, max_tokens=max_tokens,
                    )
                    if model:
                        kwargs["model"] = model
                    for chunk in adapter.analyze(img_b64, OCR_READER_PROMPT, **kwargs):
                        raw += chunk
                    if raw.strip():
                        log.info("[ocr-reader] OCR via %s: %d chars", fb_provider, len(raw))
                finally:
                    config._data.setdefault("vision", {})["provider"] = original_provider
            except Exception as exc:
                log.warning("[ocr-reader] %s failed: %s", fb_provider, exc)

    # LAST RESORT: llamacpp (slow 40-90s but works fully offline)
    if not raw.strip():
        try:
            from ..ai.provider_llamacpp import get_llamacpp_provider
            lcpp = get_llamacpp_provider()
            if lcpp is not None and lcpp.vision_ready:
                log.info("[ocr-reader] LAST RESORT: using llama-cpp-python vision (slow, no network)")
                raw = lcpp.generate(
                    OCR_READER_PROMPT,
                    image_base64=img_b64,
                    system_prompt="You are an OCR assistant. Extract ALL visible text from the image.",
                    max_tokens=max_tokens,
                    temperature=0.0,
                )
        except Exception as lcpp_exc:
            log.warning("[ocr-reader] llamacpp vision failed: %s", lcpp_exc)

    text = raw.strip()
    if not text:
        return None

    log.info("[ocr-reader] Extracted %d chars of text", len(text))
    return text


# ===================================================================== #
# Audio feedback
# ===================================================================== #

def beep_start() -> None:
    """Short beep to indicate AI request started."""
    threading.Thread(target=_do_beep, args=(800, 80), daemon=True).start()


def beep_result() -> None:
    """Short beep to indicate result received."""
    threading.Thread(target=_do_beep, args=(1200, 100), daemon=True).start()


def _do_beep(freq: int, duration_ms: int) -> None:
    """Platform-specific beep."""
    try:
        if sys.platform == "win32":
            import ctypes
            ctypes.windll.kernel32.Beep(freq, duration_ms)  # type: ignore[attr-defined]
        elif sys.platform == "darwin":
            # macOS: use NSBeep (simple system sound)
            import subprocess
            subprocess.run(
                ["afplay", "/System/Library/Sounds/Tink.aiff"],
                timeout=2, capture_output=True,
            )
    except Exception:
        pass
