"""PageSolver — single-hotkey full-page exam question solver.

Press once  → scrolls to top, then scans the page top-to-bottom,
              answering ALL questions automatically.
Press again → cancel.

Supported question types (detected from screenshot):
  - Multiple choice (A/B/C/D)  → presses the correct letter key
  - Short text / fill-in       → types the answer
  - Essay / composition        → writes a full text via text AI

Hotkey: configured via config.page_solver_hotkey (default: ctrl+f9)
"""

from __future__ import annotations

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

if TYPE_CHECKING:
    from ..config import Config
    from ..typing.engine import TypingEngine

log = logging.getLogger(__name__)

# Pre-compiled regex for parsing JSON arrays from AI response.
# Greedy match: from first "[{" to last "}]" to avoid stopping early
# on brackets inside question text (e.g. [ANSWERED: B]).
_JSON_ARRAY_RE = re.compile(r"\[[\s\S]*\]", re.DOTALL)

# Cached RapidOCR availability — tested once per process, not every step.
# None = not yet tested, True = works, False = failed (DLL error etc.)
_RAPIDOCR_OK: bool | None = None

# ------------------------------------------------------------------ #
# Two-step AI prompts
#
# STEP 1: vision model (minicpm-v) — reads/transcribes the questions
# STEP 2: text model (qwen2.5:7b)  — determines correct answers
#
# Two-step is more accurate than asking the vision model to reason.
# ------------------------------------------------------------------ #

_VISION_EXTRACT_PROMPT = """\
Look at this screenshot and list ALL text that requires an answer or response:
- Questions (multiple choice, short answer, essay)
- Form fields (empty input boxes, text areas)
- Tasks or instructions asking for input
- Quizzes, tests, or any interactive content

For each item:
1. Write its number and full text with all answer options exactly as shown.
2. If the item already has an answer/input (filled radio, typed text, highlighted option) add [ANSWERED].
   If empty/unanswered, add [UNANSWERED].

Format as a numbered list. Transcribe text accurately (Chinese, Russian, etc.).
Include ALL visible items top to bottom — even if just one word or phrase.\
"""

_TEXT_ANSWER_PROMPT = """\
You are an assistant that answers questions, quizzes, tests, and form fields. Determine the correct answer for each UNANSWERED item. Skip items marked [ANSWERED].

CONTENT FROM SCREENSHOT:
{questions}

Reply with ONLY a JSON array, no explanation:
[
  {{"q": 1, "type": "skip", "value": ""}},
  {{"q": 2, "type": "key", "value": "b"}},
  {{"q": 3, "type": "text", "value": "answer text"}},
  {{"q": 4, "type": "essay", "value": "full topic + all requirements"}}
]

Types:
- "skip"  → already answered — do NOT interact
- "key"   → multiple choice — value = single lowercase letter (a/b/c/d)
- "text"  → short answer — value = concise direct answer, 1 sentence max, no preamble
- "essay" → writing/composition — value = the FULL topic text including ALL requirements (word count, format, etc.)

IMPORTANT:
- For "text": answer directly and concisely. No explanations, no restating the question.
- For "essay": include ALL visible requirements in value (word count, structure requirements, etc.)
- Always respond in the same language as the question.
- If there is truly nothing to answer, return an empty array: []\
"""

# ------------------------------------------------------------------ #
# OCR pipeline — three tiers by speed:
#   1. winsdk  (0.1s)  — requires: python -m pip install winsdk + VS Build Tools
#   2. rapidocr (0.5s) — requires: python -m pip install rapidocr-onnxruntime
#   3. qwen-vl-max via proxy (~8s) — always available, no install needed
# ------------------------------------------------------------------ #


def _windows_ocr_sync(img_bytes: bytes) -> str | None:  # pragma: no cover
    """Run Windows.Media.Ocr on PNG image bytes. Returns text or None."""
    try:
        import asyncio

        from winsdk.windows.globalization import Language
        from winsdk.windows.graphics.imaging import (
            BitmapAlphaMode,
            BitmapDecoder,
            BitmapPixelFormat,
        )
        from winsdk.windows.media.ocr import OcrEngine
        from winsdk.windows.storage.streams import DataWriter, InMemoryRandomAccessStream

        async def _run() -> str:
            stream = InMemoryRandomAccessStream()
            writer = DataWriter(stream)
            writer.write_bytes(img_bytes)
            await writer.store_async()
            await writer.flush_async()
            stream.seek(0)

            decoder = await BitmapDecoder.create_async(stream)
            sw_bmp = await decoder.get_software_bitmap_async(
                BitmapPixelFormat.BGRA8, BitmapAlphaMode.PREMULTIPLIED
            )

            lang = Language("zh-Hans")
            engine = OcrEngine.try_create_from_language(lang)
            if engine is None:
                engine = OcrEngine.try_create_from_user_profile_languages()
            if engine is None:
                return ""

            result = await engine.recognize_async(sw_bmp)
            # Use line-by-line text for better structure
            lines = [ln.text for ln in result.lines]
            return "\n".join(lines)

        return asyncio.run(_run())
    except Exception as exc:
        log.debug("[page] Windows OCR unavailable: %s", exc)
        return None


def _rapidocr_sync(img_bytes: bytes) -> str | None:  # pragma: no cover
    """Run RapidOCR on PNG/JPEG image bytes. Returns text or None.

    Install: python -m pip install rapidocr-onnxruntime
    Also requires: Visual C++ Redistributable 2015-2022 (https://aka.ms/vs/17/release/vc_redist.x64.exe)
    Supports Chinese, Russian, English. ~0.5s per screenshot.
    """
    global _RAPIDOCR_OK
    if _RAPIDOCR_OK is False:
        return None  # already failed — skip without retrying
    try:
        import io

        import numpy as np
        from PIL import Image
        from rapidocr_onnxruntime import RapidOCR

        engine = RapidOCR()
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        arr = np.array(img)
        result, _ = engine(arr)
        _RAPIDOCR_OK = True
        if not result:
            return None
        lines = [item[1] for item in result if item and len(item) > 1]
        text = "\n".join(lines)
        log.debug("[page] RapidOCR: %d chars", len(text))
        return text if len(text.strip()) > 20 else None
    except Exception as exc:
        if _RAPIDOCR_OK is None:
            log.warning("[page] RapidOCR unavailable (will skip): %s", exc)
            _RAPIDOCR_OK = False
        return None


def _fix_ocr_text(text: str) -> str:
    """Fix known Windows OCR errors in multiple-choice exam context."""
    import re

    # "B." is often read as "&" — fix in option context: A. ... & ... C.
    text = re.sub(r"(?<!\w)&(?=\s)", "B.", text)
    # Chinese period → dot
    text = text.replace("\u3002", ".")
    # "0：" at end of lines that look like fill-in blanks
    return text


# ------------------------------------------------------------------ #
# Win32 key sending — reuse from auto_answerer
# ------------------------------------------------------------------ #

try:
    from .auto_answerer import _VK_MAP, _send_vk
except ImportError:  # pragma: no cover

    def _send_vk(vk_code: int) -> None:  # type: ignore[misc]
        pass

    _VK_MAP = {}  # type: ignore[assignment]

_VK_RETURN = 0x0D
_VK_TAB = 0x09
_VK_NEXT = 0x22  # Page Down
_VK_HOME = 0x24  # Home (scroll to top)


# ------------------------------------------------------------------ #
# PageSolver
# ------------------------------------------------------------------ #


class PageSolver:
    """Scans full page and answers all exam questions automatically."""

    def __init__(self, config: Config, typing_engine: TypingEngine) -> None:
        self._config = config
        self._typing_engine = typing_engine
        self._cancel = threading.Event()
        self._running = False

    # ------------------------------------------------------------------ #
    # Hotkey entry point
    # ------------------------------------------------------------------ #

    def on_hotkey(self) -> None:
        """Start solving or cancel if already running."""
        if self._running:
            log.info("[page] Cancel requested by user")
            self._cancel.set()
            return
        self._cancel.clear()
        log.info("[page] Starting full-page solver")
        threading.Thread(target=self._solve_all, daemon=True, name="page-solver").start()

    # ------------------------------------------------------------------ #
    # Main solving loop
    # ------------------------------------------------------------------ #

    def _solve_all(self) -> None:
        self._running = True
        advance_delay = float(self._config.get("page_solver.advance_delay", 0.3))

        try:
            # Wait for NetworkProtector to be ready
            try:
                from ..utils.network_protect import wait_for_network
                if not wait_for_network(timeout=0):
                    log.info("[page] Network protection still starting, waiting...")
                    wait_for_network(timeout=30)
            except ImportError:
                pass

            # Scroll to top of page first
            _send_vk(_VK_HOME)
            time.sleep(0.5)

            prev_hash = ""
            hash_stuck = 0   # consecutive steps where page hash didn't change
            empty_ai = 0     # consecutive steps where AI returned no unanswered items

            for step in range(60):  # safety limit: max 60 scroll steps
                if self._cancel.is_set():
                    log.info("[page] Cancelled at step %d", step)
                    break

                # Take screenshot
                img_b64 = self._capture()
                if not img_b64:
                    break

                # Detect stuck (page not scrolling anymore)
                curr_hash = hashlib.md5(img_b64[:2000].encode()).hexdigest()
                if curr_hash == prev_hash:
                    hash_stuck += 1
                    if hash_stuck >= 3:
                        log.info("[page] Page stopped changing — all done")
                        break
                else:
                    hash_stuck = 0
                prev_hash = curr_hash

                # Ask vision AI for all visible questions
                answers = self._analyze(img_b64)
                log.info("[page] Step %d: %d items from AI", step + 1, len(answers))

                if not answers:
                    # No unanswered questions on this screen.
                    # Use a SEPARATE counter so page scrolls don't reset it.
                    empty_ai += 1
                    if empty_ai >= 2:
                        log.info("[page] All questions answered — done")
                        break
                    self._scroll_down()
                    time.sleep(0.5)
                    continue

                # Got some answers — reset empty counter
                empty_ai = 0

                # All items are "done" → scroll or stop
                if all(a.get("type") == "done" for a in answers):
                    log.info("[page] AI: no more questions — scrolling to check")
                    empty_ai += 1
                    if empty_ai >= 2:
                        log.info("[page] All questions answered — done")
                        break
                    self._scroll_down()
                    time.sleep(1.0)
                    continue

                # Execute each answer
                answered_any = False
                for ans in answers:
                    if self._cancel.is_set():
                        break

                    kind = ans.get("type", "done")
                    value = str(ans.get("value", "")).strip()

                    if kind == "done":
                        break

                    if kind == "skip":
                        # Question already has an answer — do not touch it,
                        # just advance focus to the next field
                        log.info("[page] Skipping already-answered question")
                        self._press_advance()
                        time.sleep(0.2)
                        continue

                    if kind == "key" and value:
                        key = value[0].lower()
                        log.info("[page] Pressing key: %s", key.upper())
                        self._press_key(key)
                        time.sleep(0.3)
                        self._press_advance()
                        answered_any = True

                    elif kind == "text" and value:
                        log.info("[page] Typing text (%d chars)", len(value))
                        self._typing_engine.type_text(value)
                        time.sleep(0.4)
                        self._press_advance()
                        answered_any = True

                    elif kind == "essay" and value:
                        # value = essay topic; generate full essay via text AI
                        log.info("[page] Writing essay for: %r", value[:80])
                        essay = self._write_essay(value)
                        if essay:
                            self._typing_engine.type_text(essay)
                            time.sleep(0.5)
                            self._press_advance()
                            answered_any = True

                    if answered_any:
                        time.sleep(advance_delay)

                if not answered_any:
                    # Nothing answered on this screen → scroll down
                    self._scroll_down()
                    time.sleep(1.0)

        except Exception as exc:
            log.error("[page] Unexpected error: %s", exc)
        finally:
            self._running = False
            log.info("[page] Page solver finished")

    # ------------------------------------------------------------------ #
    # Screenshot helper
    # ------------------------------------------------------------------ #

    def _capture(self) -> str:
        try:
            from ..screenshot import capture_screen
            from ..utils.image_utils import image_to_base64, resize_image

            img = capture_screen()
            # Cap at 1920x1080; smaller screenshots pass through unchanged.
            img = resize_image(img)
            b64 = image_to_base64(img)
            log.debug("[page] Screenshot: %dx%d → %d chars", img.size[0], img.size[1], len(b64))
            return b64
        except Exception as exc:
            log.error("[page] Screenshot error: %s", exc)
            return ""

    # ------------------------------------------------------------------ #
    # Two-step AI analysis
    # ------------------------------------------------------------------ #

    def _analyze(self, img_b64: str) -> list[dict]:
        """Two-step: vision reads questions, text AI determines answers."""
        # Step 1: vision model extracts question text
        questions_text = self._vision_extract(img_b64)
        if not questions_text or not questions_text.strip():
            log.warning("[page] Vision returned no text")
            return []

        log.info("[page] Vision read %d chars of questions: %r", len(questions_text), questions_text[:300])

        # Step 2: text model determines correct answers
        answers = self._text_answer(questions_text)
        return answers

    def _vision_extract(self, img_b64: str) -> str:
        """Step 1: extract question text from screenshot.

        Tries Windows OCR first (0.1s, ~95% accuracy).
        Falls back to minicpm-v if winsdk not installed.
        """
        import base64

        img_bytes = base64.b64decode(img_b64)

        # Tier 1: winsdk Windows OCR (~0.1s, requires VS Build Tools)
        ocr_text = _windows_ocr_sync(img_bytes)
        if ocr_text and len(ocr_text.strip()) > 20:
            fixed = _fix_ocr_text(ocr_text)
            log.info("[page] Windows OCR: %d chars (tier-1 fast)", len(fixed))
            return fixed

        # Tier 2: RapidOCR (~0.5s, pip install rapidocr-onnxruntime, no compilation)
        rapid_text = _rapidocr_sync(img_bytes)
        if rapid_text and len(rapid_text.strip()) > 20:
            fixed = _fix_ocr_text(rapid_text)
            log.info("[page] RapidOCR: %d chars (tier-2 medium)", len(fixed))
            return fixed

        # Tier 3: vision model via proxy (~8s, always available)
        log.info("[page] Local OCR unavailable — using vision model (tier-3 slow)")
        from ..ai.vision_adapter import VisionAdapter

        adapter = VisionAdapter(self._config, timeout=60.0)
        raw = ""
        try:
            for chunk in adapter.analyze(
                img_b64, _VISION_EXTRACT_PROMPT,
                cancel_event=self._cancel,
                model="qwen-vl-plus",
                temperature=0.0,
                max_tokens=2048,
            ):
                raw += chunk
        except Exception as exc:
            log.error("[page] Vision extract error: %s", exc)
        return raw

    def _text_answer(self, questions_text: str) -> list[dict]:
        """Step 2: send question text to configured AI provider, get JSON answer list."""
        from ..ai.base import AIRequest, AIResponse
        from ..ai.china_adapter import ChinaAdapter

        prompt = _TEXT_ANSWER_PROMPT.format(questions=questions_text)
        request = AIRequest(
            prompt=prompt,
            system_prompt="You are an exam assistant. Output only valid JSON arrays.",
            model=getattr(self._config, "model", None),
            temperature=0.0,  # deterministic answers for exams
        )

        try:
            resp = ChinaAdapter(self._config).ask(request, stream=False)
            assert isinstance(resp, AIResponse)
            if not resp.ok:
                log.error("[page] Text AI error: %s", resp.error)
                return []
            raw = resp.text
        except Exception as exc:
            log.error("[page] Text AI error: %s", exc)
            return []

        log.info("[page] Text AI raw: %r", raw[:400])

        # Normalize: convert {q:N, type:X, value:Y} → {type:X, value:Y}
        match = _JSON_ARRAY_RE.search(raw)
        if match:
            try:
                items = json.loads(match.group())
                if isinstance(items, list):
                    # Strip the "q" field — PageSolver only needs type+value in order
                    return [
                        {"type": a.get("type", "done"), "value": a.get("value", "")} for a in items
                    ]
            except json.JSONDecodeError:
                pass

        log.warning("[page] Cannot parse text AI response: %r", raw[:200])
        return []

    # ------------------------------------------------------------------ #
    # Text AI — write essay
    # ------------------------------------------------------------------ #

    def _write_essay(self, topic: str) -> str:
        """Use configured AI provider to write an essay strictly following all requirements."""
        from ..ai.base import AIRequest, AIResponse
        from ..ai.china_adapter import ChinaAdapter

        prompt = (
            "Write a complete essay strictly following ALL requirements listed below. "
            "Do not add anything beyond what is required. "
            "Output the essay text only — no preamble, no meta-commentary, no explanations. "
            "Respond in the same language as the topic.\n\n"
            f"Requirements and topic:\n{topic}"
        )
        request = AIRequest(
            prompt=prompt,
            system_prompt=(
                "You are an expert essay writer. "
                "Write essays strictly following all given requirements. "
                "Output only the essay text, nothing else."
            ),
            model=getattr(self._config, "model", None),
            max_tokens=4096,
        )

        try:
            resp = ChinaAdapter(self._config).ask(request, stream=False)
            assert isinstance(resp, AIResponse)
            if not resp.ok:
                log.error("[page] Essay generation error: %s", resp.error)
                return ""
            essay = resp.text.strip()
            log.info("[page] Essay written: %d chars", len(essay))
            return essay
        except Exception as exc:
            log.error("[page] Essay generation error: %s", exc)
            return ""

    # ------------------------------------------------------------------ #
    # Key helpers
    # ------------------------------------------------------------------ #

    def _press_key(self, key: str) -> None:
        vk = _VK_MAP.get(key.lower())
        if not vk:
            log.warning("[page] Unknown key: %r", key)
            return
        if sys.platform == "win32":
            _send_vk(vk)
        else:
            try:
                import pyautogui

                pyautogui.press(key)
            except Exception as exc:
                log.warning("[page] pyautogui error: %s", exc)

    def _press_advance(self) -> None:
        advance = str(self._config.get("page_solver.advance_key", "tab")).lower()
        vk = _VK_RETURN if advance == "enter" else _VK_TAB
        if sys.platform == "win32":
            _send_vk(vk)
        else:
            try:
                import pyautogui

                pyautogui.press("enter" if advance == "enter" else "tab")
            except Exception as exc:
                log.warning("[page] advance key error: %s", exc)

    def _scroll_down(self) -> None:
        if sys.platform == "win32":
            _send_vk(_VK_NEXT)
        else:
            try:
                import pyautogui

                pyautogui.press("pagedown")
            except Exception as exc:
                log.warning("[page] scroll error: %s", exc)
        log.debug("[page] Scrolled down (Page Down)")
