"""上传图片 → OCR → 结构化提取 → 生成草稿题目。作为 BackgroundTask 运行。

注意：BackgroundTask 在请求返回后执行，需要自己开新的 DB session。
"""
import json
import logging

from pydantic import ValidationError
from sqlalchemy import select

from app.db import SessionLocal
from app.models.image import Image
from app.models.job import JobStatus, ProcessingJob
from app.models.question import Question, QuestionType
from app.models.user import UserSettings
from app.schemas.question import QuestionCreate
from app.services.llm.client import LLMClient, LLMConfigError
from app.services.llm.prompts import EXTRACT_SYSTEM, EXTRACT_USER_TEMPLATE
from app.services.ocr.pipeline import run_ocr
from app.services.storage import download_image_async

logger = logging.getLogger(__name__)


def _parse_extract_json(raw: str) -> dict:
    """从模型输出里解析 JSON（容忍代码块包裹）。"""
    text = raw.strip()
    if text.startswith("```"):
        text = text.split("```", 2)[1]
        if text.startswith("json"):
            text = text[4:]
    return json.loads(text.strip())


async def run_extract_job(job_id: str, user_id: str, object_key: str, mime: str):
    async with SessionLocal() as db:
        job = await db.get(ProcessingJob, job_id)
        if job is None:
            logger.error("job %s 不存在", job_id)
            return
        try:
            # 加载用户 AI 配置
            settings_row = await db.scalar(
                select(UserSettings).where(UserSettings.user_id == user_id)
            )
            llm = LLMClient(settings_row)

            # 1. OCR
            job.status = JobStatus.ocr_running
            await db.commit()
            image_bytes = await download_image_async(object_key)
            ocr = await run_ocr(image_bytes, mime, llm)

            # 2. 结构化提取
            job.status = JobStatus.ai_running
            job.engine_used = ocr.engine
            await db.commit()
            extract = await llm.chat_text(
                EXTRACT_SYSTEM,
                EXTRACT_USER_TEMPLATE.format(content=ocr.markdown),
                json_mode=True,
            )
            questions = _build_questions(user_id, extract.content, ocr, job.image_id)

            # 缓存原始 OCR 文本，便于日后追溯或重新拆分
            if job.image_id:
                image = await db.get(Image, job.image_id)
                if image is not None:
                    image.ocr_markdown = ocr.markdown

            db.add_all(questions)
            await db.flush()
            job.result_question_ids = [q.id for q in questions]
            job.status = JobStatus.done
            await db.commit()
            logger.info(
                "job %s 完成，engine=%s，拆出 %d 条笔记",
                job_id,
                ocr.engine,
                len(questions),
            )
        except LLMConfigError as e:
            await _fail(db, job, f"AI 未配置：{e}")
        except Exception as e:  # noqa: BLE001
            logger.exception("job %s 失败", job_id)
            await _fail(db, job, str(e))


def _build_questions(
    user_id: str, extract_raw: str, ocr, image_id: str | None
) -> list[Question]:
    """把提取结果拆成多条笔记。

    一张照片常含多道题，每道题存成独立笔记。整体解析失败时不丢内容，
    降级为一条 unclassified 笔记（原文进题干），用户可自行整理。
    """
    try:
        data = _parse_extract_json(extract_raw)
        raw_items = data.get("questions")
        if not isinstance(raw_items, list) or not raw_items:
            raise ValueError("提取结果里没有 questions 数组")
    except (json.JSONDecodeError, KeyError, ValueError) as e:
        logger.warning("提取结果解析失败，降级为一条未分类笔记：%s", e)
        return [_fallback_question(user_id, ocr, image_id)]

    questions: list[Question] = []
    for idx, item in enumerate(raw_items):
        try:
            questions.append(_build_one(user_id, item, ocr, image_id))
        except (KeyError, ValueError, ValidationError, TypeError) as e:
            # 单条坏掉不影响其他条：这条退化为未分类，原文留在题干里
            logger.warning("第 %d 条解析失败，存为未分类：%s", idx + 1, e)
            stem = ""
            if isinstance(item, dict):
                stem = str(item.get("stem_markdown") or "")
            questions.append(
                _fallback_question(user_id, ocr, image_id, stem_override=stem or None)
            )

    return questions or [_fallback_question(user_id, ocr, image_id)]


def _build_one(user_id: str, item: dict, ocr, image_id: str | None) -> Question:
    parsed = QuestionCreate(
        # 模型没给或给了不认识的题型，就当未分类，别硬猜
        type=_coerce_type(item.get("type")),
        stem_markdown=item.get("stem_markdown") or ocr.markdown,
        options=item.get("options"),
        correct_answer=item.get("correct_answer"),
        difficulty=item.get("difficulty"),
    )
    return Question(
        user_id=user_id,
        type=parsed.type,
        stem_markdown=parsed.stem_markdown,
        options=[o.model_dump() for o in parsed.options] if parsed.options else None,
        correct_answer=parsed.correct_answer,
        difficulty=parsed.difficulty,
        source_image_id=image_id,
        ocr_engine=ocr.engine,
    )


def _coerce_type(raw) -> QuestionType:
    try:
        return QuestionType(raw)
    except ValueError:
        return QuestionType.unclassified


def _fallback_question(
    user_id: str, ocr, image_id: str | None, stem_override: str | None = None
) -> Question:
    return Question(
        user_id=user_id,
        type=QuestionType.unclassified,
        stem_markdown=stem_override or ocr.markdown,
        source_image_id=image_id,
        ocr_engine=ocr.engine,
    )


async def _fail(db, job: ProcessingJob, msg: str):
    job.status = JobStatus.failed
    job.error_message = msg[:1000]
    await db.commit()