#!/usr/bin/env python3
"""Asynchronously generate multiple solutions for a problem and store into SQLite."""

from __future__ import annotations

import argparse
import json
import os
import re
import sqlite3
import time
from dataclasses import dataclass
from typing import Any

import requests

RETRYABLE_HTTP_CODES = {500, 502, 503, 504}


@dataclass
class Problem:
    id: int
    title: str
    statement_md: str
    difficulty: int
    source: str
    sample_input: str
    sample_output: str


def now_sec() -> int:
    return int(time.time())


def extract_json_object(text: str) -> dict[str, Any] | None:
    raw = text.strip()
    if raw.startswith("```"):
        raw = re.sub(r"^```[a-zA-Z0-9_-]*", "", raw).strip()
        raw = raw.removesuffix("```").strip()
    try:
        obj = json.loads(raw)
        if isinstance(obj, dict):
            return obj
    except json.JSONDecodeError:
        pass

    match = re.search(r"\{[\s\S]*\}", text)
    if not match:
        return None
    try:
        obj = json.loads(match.group(0))
        return obj if isinstance(obj, dict) else None
    except json.JSONDecodeError:
        return None


def llm_request(prompt: str, timeout: int, retries: int, sleep_sec: float) -> str:
    url = os.getenv("OI_LLM_API_URL", "").strip()
    api_key = os.getenv("OI_LLM_API_KEY", "").strip()
    model = os.getenv("OI_LLM_MODEL", "qwen3-max").strip()
    if not url:
        raise RuntimeError("missing OI_LLM_API_URL")

    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    body = {
        "model": model,
        "stream": False,
        "temperature": 0.3,
        "messages": [
            {
                "role": "system",
                "content": "你是资深 OI/CSP 教练。严格输出 JSON，不要输出任何额外文本。",
            },
            {"role": "user", "content": prompt},
        ],
    }

    last_error: Exception | None = None
    for i in range(1, retries + 1):
        try:
            resp = requests.post(url, headers=headers, json=body, timeout=timeout)
        except requests.RequestException as exc:
            last_error = exc
            if i < retries:
                time.sleep(sleep_sec * i)
                continue
            raise RuntimeError(f"llm request failed: {exc}") from exc

        if resp.status_code in RETRYABLE_HTTP_CODES:
            if i < retries:
                time.sleep(sleep_sec * i)
                continue
            raise RuntimeError(f"llm retry exhausted: HTTP {resp.status_code}")

        if resp.status_code >= 400:
            raise RuntimeError(f"llm request failed: HTTP {resp.status_code}: {resp.text[:300]}")

        payload = resp.json()
        choices = payload.get("choices") or []
        if not choices:
            raise RuntimeError("llm response missing choices")
        content = ((choices[0] or {}).get("message") or {}).get("content")
        if not content:
            raise RuntimeError("llm response missing content")
        return str(content)

    if last_error:
        raise RuntimeError(f"llm request failed: {last_error}") from last_error
    raise RuntimeError("llm request failed")


def fallback_solutions(max_solutions: int) -> list[dict[str, Any]]:
    base = [
        {
            "title": "解法一：直接模拟/枚举",
            "idea_md": "按题意拆分步骤，先写可过样例的直观解法，再补边界处理。",
            "explanation_md": "适用于数据范围较小或规则清晰的题。",
            "complexity": "时间复杂度依题而定，通常 O(n)~O(n^2)",
            "code_cpp": "// TODO: 请根据题意补全\n#include <bits/stdc++.h>\nusing namespace std;\nint main(){ios::sync_with_stdio(false);cin.tie(nullptr);return 0;}\n",
            "tags": ["simulation", "implementation"],
        },
        {
            "title": "解法二：优化思路（前缀/贪心/DP 视题而定）",
            "idea_md": "分析状态与重复计算，尝试用前缀和、贪心或动态规划优化。",
            "explanation_md": "比直接模拟更稳定，通常能覆盖更大数据规模。",
            "complexity": "通常优于朴素解法",
            "code_cpp": "// TODO: 请根据题意补全\n#include <bits/stdc++.h>\nusing namespace std;\nint main(){ios::sync_with_stdio(false);cin.tie(nullptr);return 0;}\n",
            "tags": ["optimization", "dp"],
        },
    ]
    return base[: max(1, max_solutions)]


def load_problem(conn: sqlite3.Connection, problem_id: int) -> Problem:
    cur = conn.execute(
        "SELECT id,title,statement_md,difficulty,source,sample_input,sample_output FROM problems WHERE id=?",
        (problem_id,),
    )
    row = cur.fetchone()
    if row is None:
        raise RuntimeError(f"problem not found: {problem_id}")
    return Problem(
        id=int(row[0]),
        title=str(row[1] or ""),
        statement_md=str(row[2] or ""),
        difficulty=int(row[3] or 1),
        source=str(row[4] or ""),
        sample_input=str(row[5] or ""),
        sample_output=str(row[6] or ""),
    )


def update_job(conn: sqlite3.Connection, job_id: int, **fields: Any) -> None:
    if not fields:
        return
    keys = []
    vals: list[Any] = []
    for k, v in fields.items():
        keys.append(f"{k}=?")
        vals.append(v)
    vals.append(job_id)
    conn.execute(
        f"UPDATE problem_solution_jobs SET {', '.join(keys)} WHERE id=?",
        tuple(vals),
    )
    conn.commit()


def store_solutions(conn: sqlite3.Connection, problem_id: int, rows: list[dict[str, Any]], source: str) -> int:
    ts = now_sec()
    conn.execute("DELETE FROM problem_solutions WHERE problem_id=?", (problem_id,))
    saved = 0
    seen_titles: set[str] = set()
    for idx, row in enumerate(rows, start=1):
        title = str(row.get("title") or f"解法 {idx}").strip()
        if title in seen_titles:
            continue
        seen_titles.add(title)

        idea_md = str(row.get("idea_md") or "").strip()
        explanation_md = str(row.get("explanation_md") or "").strip()
        code_cpp = str(row.get("code_cpp") or "").strip()
        complexity = str(row.get("complexity") or "").strip()
        tags = row.get("tags") if isinstance(row.get("tags"), list) else []

        conn.execute(
            """
            INSERT INTO problem_solutions(
              problem_id,variant,title,idea_md,explanation_md,code_cpp,complexity,tags_json,source,created_at,updated_at
            ) VALUES(?,?,?,?,?,?,?,?,?,?,?)
            """,
            (
                problem_id,
                idx,
                title,
                idea_md,
                explanation_md,
                code_cpp,
                complexity,
                json.dumps(tags, ensure_ascii=False),
                source,
                ts,
                ts,
            ),
        )
        saved += 1
    conn.commit()
    return saved


def main() -> int:
    parser = argparse.ArgumentParser(description="Generate multi-solution explanations")
    parser.add_argument("--db-path", required=True)
    parser.add_argument("--problem-id", type=int, required=True)
    parser.add_argument("--job-id", type=int, required=True)
    parser.add_argument("--max-solutions", type=int, default=3)
    parser.add_argument("--timeout", type=int, default=90)
    parser.add_argument("--retries", type=int, default=4)
    parser.add_argument("--retry-sleep-sec", type=float, default=1.5)
    args = parser.parse_args()

    conn = sqlite3.connect(args.db_path)
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA busy_timeout=5000")

    ts = now_sec()
    update_job(
        conn,
        args.job_id,
        status="running",
        progress=1,
        message="starting",
        started_at=ts,
        updated_at=ts,
    )

    try:
        problem = load_problem(conn, args.problem_id)

        prompt = f"""
请为下面这道 CSP 题生成 {max(1, min(5, args.max_solutions))} 种不同思路的题解（可从不同角度切入，例如模拟/贪心/DP/数据结构），并给出 C++ 参考代码。

输出 JSON，格式固定：
{{
  "solutions": [
    {{
      "title": "解法标题",
      "idea_md": "思路要点（Markdown）",
      "explanation_md": "详细讲解（Markdown）",
      "complexity": "时间/空间复杂度",
      "code_cpp": "完整 C++17 代码",
      "tags": ["标签1","标签2"]
    }}
  ]
}}

题目：{problem.title}
难度：{problem.difficulty}
来源：{problem.source}
题面：
{problem.statement_md[:12000]}
样例输入：
{problem.sample_input[:1200]}
样例输出：
{problem.sample_output[:1200]}
""".strip()

        update_job(conn, args.job_id, progress=25, message="requesting llm", updated_at=now_sec())

        source = "fallback"
        solutions: list[dict[str, Any]]
        try:
            content = llm_request(
                prompt,
                timeout=args.timeout,
                retries=args.retries,
                sleep_sec=args.retry_sleep_sec,
            )
            obj = extract_json_object(content)
            raw = obj.get("solutions") if isinstance(obj, dict) else None
            if not isinstance(raw, list) or len(raw) == 0:
                raise RuntimeError("llm response missing solutions array")
            solutions = [x for x in raw if isinstance(x, dict)]
            if not solutions:
                raise RuntimeError("llm response has empty valid solutions")
            source = "llm"
        except Exception:
            solutions = fallback_solutions(args.max_solutions)

        solutions = solutions[: max(1, min(5, args.max_solutions))]

        update_job(conn, args.job_id, progress=70, message="writing solutions", updated_at=now_sec())
        saved = store_solutions(conn, args.problem_id, solutions, source)

        update_job(
            conn,
            args.job_id,
            status="completed",
            progress=100,
            message=f"completed: {saved} solutions ({source})",
            finished_at=now_sec(),
            updated_at=now_sec(),
        )
        conn.close()
        return 0
    except Exception as exc:
        update_job(
            conn,
            args.job_id,
            status="failed",
            progress=100,
            message=f"failed: {str(exc)[:400]}",
            finished_at=now_sec(),
            updated_at=now_sec(),
        )
        conn.close()
        return 1


if __name__ == "__main__":
    raise SystemExit(main())