import html
import json
import logging
import os
import random
import re
from dataclasses import dataclass, field
from typing import Any

import psycopg
from psycopg.rows import dict_row
from telegram import BotCommand, InlineKeyboardButton, InlineKeyboardMarkup, Update
from telegram.constants import ChatAction, ParseMode
from telegram.ext import Application, CallbackQueryHandler, CommandHandler, ContextTypes, MessageHandler, filters


logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("quiz-telegram-bot")
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)

TELEGRAM_BOT_TOKEN = os.environ["TELEGRAM_BOT_TOKEN"]
DATABASE_URL = os.environ["DATABASE_URL"]
DEFAULT_QUIZ_SIZE = int(os.getenv("DEFAULT_QUIZ_SIZE", "20"))
MAX_QUIZ_SIZE = int(os.getenv("MAX_QUIZ_SIZE", "50"))
PUBLIC_APP_URL = os.getenv("PUBLIC_APP_URL", "https://pedshub.com").rstrip("/")
TELEGRAM_MESSAGE_LIMIT = 4096

LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
NUMBER_RE = re.compile(r"^\s*(\d{1,3})\s*$")


@dataclass
class QuizState:
    questions: list[dict[str, Any]]
    mode: str = "study"
    index: int = 0
    score: int = 0
    answers: list[tuple[int, str, str, bool]] = field(default_factory=list)


active_quizzes: dict[int, QuizState] = {}


def db_query(sql: str, params: tuple[Any, ...] = ()) -> list[dict[str, Any]]:
    with psycopg.connect(DATABASE_URL, row_factory=dict_row) as conn:
        with conn.cursor() as cur:
            cur.execute(sql, params)
            return list(cur.fetchall())


def normalize_options(value: Any) -> list[str]:
    if value is None:
        return []
    if isinstance(value, list):
        return [str(item) for item in value]
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            return []
        if isinstance(parsed, list):
            return [str(item) for item in parsed]
    return []


def answer_index(question: dict[str, Any], options: list[str]) -> int | None:
    raw = str(question.get("correct_answer") or "").strip()
    if not raw:
        return None
    upper = raw.upper()
    if len(upper) == 1 and upper in LETTERS:
        idx = LETTERS.index(upper)
        return idx if idx < len(options) else None
    if raw.isdigit():
        idx = int(raw) - 1
        return idx if 0 <= idx < len(options) else None
    for idx, option in enumerate(options):
        if option.strip().lower() == raw.lower():
            return idx
    return None


def question_base_sql(where: str = "") -> str:
    return f"""
        select q.id, q.question_text, q.options, q.correct_answer, q.explanation, qc.name as category
        from questions q
        left join question_categories qc on qc.id = q.question_category_id
        where q.question_type = 'mcq'
          and q.options is not null
          and q.is_shared = 1
          {where}
        order by random()
        limit %s
    """


def valid_questions(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    clean = []
    for row in rows:
        options = normalize_options(row.get("options"))
        if len(options) < 2:
            continue
        idx = answer_index(row, options)
        if idx is None:
            continue
        row["options"] = options
        row["correct_index"] = idx
        clean.append(row)
    return clean


def random_questions(limit: int) -> list[dict[str, Any]]:
    rows = db_query(question_base_sql(), (limit * 2,))
    return valid_questions(rows)[:limit]


def search_questions(term: str, limit: int) -> list[dict[str, Any]]:
    pattern = f"%{term}%"
    rows = db_query(
        question_base_sql("""
          and (
            q.question_text ilike %s
            or q.explanation ilike %s
            or qc.name ilike %s
            or exists (
              select 1
              from question_tag_links qtl
              join question_tags qt on qt.id = qtl.tag_id
              where qtl.question_id = q.id and qt.name ilike %s
            )
          )
        """),
        (pattern, pattern, pattern, pattern, limit * 2),
    )
    return valid_questions(rows)[:limit]


def category_questions(category_id: int, limit: int) -> list[dict[str, Any]]:
    rows = db_query(question_base_sql("and q.question_category_id = %s"), (category_id, limit * 2))
    return valid_questions(rows)[:limit]


def tag_questions(tag_id: int, limit: int) -> list[dict[str, Any]]:
    rows = db_query(
        question_base_sql("and exists (select 1 from question_tag_links qtl where qtl.question_id = q.id and qtl.tag_id = %s)"),
        (tag_id, limit * 2),
    )
    return valid_questions(rows)[:limit]


def list_categories() -> list[dict[str, Any]]:
    return db_query(
        """
        select qc.id, qc.name, count(q.id)::int as count
        from question_categories qc
        join questions q on q.question_category_id = qc.id and q.question_type = 'mcq' and q.is_shared = 1
        group by qc.id, qc.name
        having count(q.id) > 0
        order by qc.name
        """
    )


def search_tags(term: str | None = None, limit: int = 20) -> list[dict[str, Any]]:
    where = ""
    params: list[Any] = []
    if term:
        where = "where t.name ilike %s"
        params.append(f"%{term}%")
    params.append(limit)
    return db_query(
        f"""
        select t.id, t.name, t.type, count(qtl.question_id)::int as count
        from question_tags t
        join question_tag_links qtl on qtl.tag_id = t.id
        join questions q on q.id = qtl.question_id and q.question_type = 'mcq' and q.is_shared = 1
        {where}
        group by t.id, t.name, t.type
        order by count(qtl.question_id) desc, t.name
        limit %s
        """,
        tuple(params),
    )


def clamp_count(value: int | None) -> int:
    if value is None:
        return DEFAULT_QUIZ_SIZE
    return max(1, min(MAX_QUIZ_SIZE, value))


def parse_count(args: list[str]) -> int:
    for arg in args:
        if arg.isdigit():
            return clamp_count(int(arg))
    return DEFAULT_QUIZ_SIZE


def parse_mode(args: list[str]) -> str:
    lowered = {arg.lower() for arg in args}
    return "exam" if "exam" in lowered else "study"


async def send_text(update: Update, text: str, reply_markup: InlineKeyboardMarkup | None = None) -> None:
    if update.message:
        await update.message.reply_text(
            text,
            parse_mode=ParseMode.HTML,
            disable_web_page_preview=True,
            reply_markup=reply_markup,
        )


async def respond(update: Update, context: ContextTypes.DEFAULT_TYPE | None, text: str, reply_markup: InlineKeyboardMarkup | None = None) -> None:
    if update.message:
        await send_text(update, text, reply_markup)
    elif update.callback_query and update.callback_query.message:
        await update.callback_query.edit_message_text(text, parse_mode=ParseMode.HTML, reply_markup=reply_markup, disable_web_page_preview=True)
    elif context and update.effective_chat:
        await context.bot.send_message(update.effective_chat.id, text, parse_mode=ParseMode.HTML, reply_markup=reply_markup, disable_web_page_preview=True)


def main_menu() -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup([
        [InlineKeyboardButton("Random 20 study", callback_data="quiz:random:20:study")],
        [InlineKeyboardButton("Random 20 exam", callback_data="quiz:random:20:exam")],
        [InlineKeyboardButton("Categories", callback_data="list:categories:0")],
        [InlineKeyboardButton("Top keywords", callback_data="list:tags:0")],
    ])


def count_menu(kind: str, item_id: int) -> InlineKeyboardMarkup:
    buttons = []
    for count in (5, 10, 20, 30, 50):
        buttons.append([
            InlineKeyboardButton(f"{count} study", callback_data=f"start:{kind}:{item_id}:{count}:study"),
            InlineKeyboardButton(f"{count} exam", callback_data=f"start:{kind}:{item_id}:{count}:exam"),
        ])
    return InlineKeyboardMarkup(buttons)


def format_question_review(state: QuizState, question: dict[str, Any], chosen_idx: int, correct_idx: int, ok: bool) -> str:
    lines = [
        f"<b>Question {state.index + 1}/{len(state.questions)}</b>",
        html.escape(question["question_text"]),
        "",
    ]
    for idx, option in enumerate(question["options"][:8]):
        marker = ""
        if idx == correct_idx:
            marker = " correct"
        elif idx == chosen_idx:
            marker = " your answer"
        lines.append(f"{LETTERS[idx]}. {html.escape(option)}{marker}")
    lines.extend([
        "",
        "Correct." if ok else f"Incorrect. Correct answer: {LETTERS[correct_idx]}",
        "",
        f"<b>Explanation:</b> {html.escape(question.get('explanation') or 'No explanation available.')}",
    ])
    return "\n".join(lines)


def truncate_text(value: str, limit: int) -> str:
    value = (value or "").strip()
    if len(value) <= limit:
        return value
    return value[:limit].rsplit(" ", 1)[0].rstrip() + "..."


def answer_feedback_header(chosen_idx: int, correct_idx: int, ok: bool) -> str:
    chosen = LETTERS[chosen_idx] if chosen_idx < len(LETTERS) else "?"
    correct = LETTERS[correct_idx] if correct_idx < len(LETTERS) else "?"
    return "Correct." if ok else f"Incorrect. You chose {chosen}; correct answer: {correct}."


def split_plain_text(text: str, limit: int) -> list[str]:
    if len(text) <= limit:
        return [text]
    chunks = []
    remaining = text
    while remaining:
        if len(remaining) <= limit:
            chunks.append(remaining)
            break
        split_at = remaining.rfind("\n\n", 0, limit)
        if split_at < limit // 2:
            split_at = remaining.rfind("\n", 0, limit)
        if split_at < limit // 2:
            split_at = remaining.rfind(" ", 0, limit)
        if split_at < limit // 2:
            split_at = limit
        chunks.append(remaining[:split_at].rstrip())
        remaining = remaining[split_at:].lstrip()
    return chunks


async def send_answer_feedback(
    context: ContextTypes.DEFAULT_TYPE,
    chat_id: int,
    question: dict[str, Any],
    chosen_idx: int,
    correct_idx: int,
    ok: bool,
) -> None:
    header = answer_feedback_header(chosen_idx, correct_idx, ok)
    explanation = (question.get("explanation") or "No explanation available.").strip()
    first_prefix = f"{header}\n\n<b>Explanation:</b> "
    next_prefix = "<b>Explanation continued:</b> "
    first_limit = TELEGRAM_MESSAGE_LIMIT - len(first_prefix) - 200
    next_limit = TELEGRAM_MESSAGE_LIMIT - len(next_prefix) - 200
    chunks = split_plain_text(explanation, max(1000, first_limit))
    for idx, chunk in enumerate(chunks):
        prefix = first_prefix if idx == 0 else next_prefix
        if idx > 0 and len(chunk) > next_limit:
            # First chunk has a smaller budget because it includes answer feedback.
            for subchunk in split_plain_text(chunk, next_limit):
                await context.bot.send_message(chat_id, f"{next_prefix}{html.escape(subchunk)}", parse_mode=ParseMode.HTML)
            continue
        await context.bot.send_message(chat_id, f"{prefix}{html.escape(chunk)}", parse_mode=ParseMode.HTML)


def answered_question_text(state: QuizState, question: dict[str, Any], chosen_idx: int, correct_idx: int) -> str:
    lines = [
        f"<b>Question {state.index + 1}/{len(state.questions)}</b>",
        html.escape(truncate_text(question["question_text"], 1600)),
        "",
    ]
    for idx, option in enumerate(question["options"][:8]):
        marker = ""
        if idx == correct_idx:
            marker = " ✓ correct"
        elif idx == chosen_idx:
            marker = " ✗ your answer"
        lines.append(f"{LETTERS[idx]}. {html.escape(truncate_text(option, 420))}{marker}")
    lines.extend([
        "",
        f"Category: {html.escape(str(question.get('category') or 'Uncategorized'))}",
    ])
    return "\n".join(lines)


async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    await send_text(update, "\n".join([
        "<b>PedQuiz bot</b>",
        "Send <code>20</code> for a random 20-question study quiz.",
        "Use /random 20 exam for exam mode.",
        "Use /categories to choose from PREP/category buckets.",
        "Use /keywords fever to search generated keywords/subjects.",
        "Use /search sepsis 20 to quiz by text search.",
        "Use /stop to end the current quiz.",
        "",
        "Study mode shows each answer immediately. Exam mode shows answers at the end.",
    ]), main_menu())


async def random_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    await start_quiz(update, random_questions(parse_count(context.args)), parse_mode(context.args))


async def search_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not context.args:
        await send_text(update, "Usage: /search <topic> [number] [study|exam]")
        return
    count = parse_count(context.args)
    mode = parse_mode(context.args)
    term = " ".join(arg for arg in context.args if not arg.isdigit() and arg.lower() not in {"study", "exam"}).strip()
    if not term:
        await send_text(update, "Usage: /search <topic> [number] [study|exam]")
        return
    if update.message:
        await update.message.chat.send_action(ChatAction.TYPING)
    await start_quiz(update, search_questions(term, count), mode, label=f"Search: {term}")


async def categories_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    await show_categories(update, page=0)


async def keywords_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    term = " ".join(context.args).strip() if context.args else None
    rows = search_tags(term, 30)
    if not rows:
        await send_text(update, "No keywords found.")
        return
    buttons = [
        [InlineKeyboardButton(f"{row['name']} ({row['type']}, {row['count']})", callback_data=f"pick:tag:{row['id']}")]
        for row in rows[:20]
    ]
    await send_text(update, "Choose a keyword/subject, then choose quiz size:", InlineKeyboardMarkup(buttons))


async def stop_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.effective_chat:
        active_quizzes.pop(update.effective_chat.id, None)
    await send_text(update, "Quiz stopped.", main_menu())


async def show_categories(update: Update, page: int) -> None:
    rows = list_categories()
    per_page = 10
    start = page * per_page
    chunk = rows[start:start + per_page]
    if not chunk:
        await send_text(update, "No categories found.")
        return
    buttons = [
        [InlineKeyboardButton(f"{row['name']} ({row['count']})", callback_data=f"pick:category:{row['id']}")]
        for row in chunk
    ]
    nav = []
    if page > 0:
        nav.append(InlineKeyboardButton("Prev", callback_data=f"list:categories:{page - 1}"))
    if start + per_page < len(rows):
        nav.append(InlineKeyboardButton("Next", callback_data=f"list:categories:{page + 1}"))
    if nav:
        buttons.append(nav)
    text = "Choose a category, then choose quiz size and mode."
    if update.callback_query:
        await update.callback_query.edit_message_text(text, reply_markup=InlineKeyboardMarkup(buttons))
    else:
        await send_text(update, text, InlineKeyboardMarkup(buttons))


async def show_tags(update: Update, page: int) -> None:
    rows = search_tags(limit=80)
    per_page = 10
    start = page * per_page
    chunk = rows[start:start + per_page]
    buttons = [
        [InlineKeyboardButton(f"{row['name']} ({row['type']}, {row['count']})", callback_data=f"pick:tag:{row['id']}")]
        for row in chunk
    ]
    nav = []
    if page > 0:
        nav.append(InlineKeyboardButton("Prev", callback_data=f"list:tags:{page - 1}"))
    if start + per_page < len(rows):
        nav.append(InlineKeyboardButton("Next", callback_data=f"list:tags:{page + 1}"))
    if nav:
        buttons.append(nav)
    text = "Choose a keyword/subject, then choose quiz size and mode."
    if update.callback_query:
        await update.callback_query.edit_message_text(text, reply_markup=InlineKeyboardMarkup(buttons))
    else:
        await send_text(update, text, InlineKeyboardMarkup(buttons))


async def start_quiz(
    update: Update,
    questions: list[dict[str, Any]],
    mode: str,
    label: str = "Random",
    context: ContextTypes.DEFAULT_TYPE | None = None,
) -> None:
    if not update.effective_chat:
        return
    if not questions:
        await respond(update, context, "No usable MCQ questions found for that selection. Try /keywords, /categories, or a broader /search term.", main_menu())
        return
    random.shuffle(questions)
    active_quizzes[update.effective_chat.id] = QuizState(questions=questions, mode=mode)
    start_text = f"{html.escape(label)} quiz started: {len(questions)} questions, {mode} mode."
    if update.message:
        await send_text(update, start_text)
        await send_current_question(update.effective_chat.id, update, None)
    elif context:
        await context.bot.send_message(update.effective_chat.id, start_text, parse_mode=ParseMode.HTML)
        await send_current_question(update.effective_chat.id, None, context)


async def send_current_question(chat_id: int, update: Update | None, context: ContextTypes.DEFAULT_TYPE | None) -> None:
    state = active_quizzes.get(chat_id)
    if not state:
        return
    question = state.questions[state.index]
    options = question["options"]
    buttons = [
        [InlineKeyboardButton(LETTERS[idx], callback_data=f"answer:{idx}")]
        for idx, option in enumerate(options[:8])
    ]
    lines = [
        f"<b>Question {state.index + 1}/{len(state.questions)}</b>",
        html.escape(truncate_text(question["question_text"], 1600)),
        "",
    ]
    for idx, option in enumerate(options[:8]):
        lines.append(f"{LETTERS[idx]}. {html.escape(truncate_text(option, 700))}")
    lines.extend([
        "",
        f"Category: {html.escape(str(question.get('category') or 'Uncategorized'))}",
    ])
    text = "\n".join(lines)
    markup = InlineKeyboardMarkup(buttons)
    if update and update.message:
        await update.message.reply_text(text, parse_mode=ParseMode.HTML, reply_markup=markup)
    elif context:
        await context.bot.send_message(chat_id, text, parse_mode=ParseMode.HTML, reply_markup=markup)


async def finish_quiz(chat_id: int, context: ContextTypes.DEFAULT_TYPE) -> None:
    state = active_quizzes.pop(chat_id, None)
    if not state:
        return
    total = len(state.questions)
    lines = [f"Quiz complete. Score: {state.score}/{total}"]
    if state.mode == "exam":
        lines.append("")
        lines.append("Answers:")
        for question_id, chosen, correct, ok in state.answers:
            marker = "OK" if ok else "MISS"
            lines.append(f"Q{question_id}: {marker}. You: {chosen}. Correct: {correct}")
    lines.append("")
    lines.append(f"Full question bank: {PUBLIC_APP_URL}")
    await context.bot.send_message(chat_id, "\n".join(lines), reply_markup=main_menu(), disable_web_page_preview=True)


async def callback_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    if not query:
        return
    await query.answer()
    data = query.data or ""
    chat_id = query.message.chat_id if query.message else update.effective_chat.id

    if data.startswith("list:categories:"):
        await show_categories(update, int(data.rsplit(":", 1)[1]))
        return
    if data.startswith("list:tags:"):
        await show_tags(update, int(data.rsplit(":", 1)[1]))
        return
    if data.startswith("quiz:random:"):
        _, _, count, mode = data.split(":", 3)
        await start_quiz(update, random_questions(clamp_count(int(count))), mode, context=context)
        return
    if data.startswith("pick:category:"):
        category_id = int(data.rsplit(":", 1)[1])
        await query.edit_message_text("How many questions?", reply_markup=count_menu("category", category_id))
        return
    if data.startswith("pick:tag:"):
        tag_id = int(data.rsplit(":", 1)[1])
        await query.edit_message_text("How many questions?", reply_markup=count_menu("tag", tag_id))
        return
    if data.startswith("start:category:"):
        _, _, category_id, count, mode = data.split(":", 4)
        await start_quiz(update, category_questions(int(category_id), clamp_count(int(count))), mode, "Category", context)
        return
    if data.startswith("start:tag:"):
        _, _, tag_id, count, mode = data.split(":", 4)
        await start_quiz(update, tag_questions(int(tag_id), clamp_count(int(count))), mode, "Keyword", context)
        return
    if data.startswith("answer:"):
        state = active_quizzes.get(chat_id)
        if not state:
            await query.edit_message_text("This quiz expired. Start a new one with /random or /categories.")
            return
        chosen_idx = int(data.split(":", 1)[1])
        question = state.questions[state.index]
        correct_idx = question["correct_index"]
        ok = chosen_idx == correct_idx
        if ok:
            state.score += 1
        chosen = LETTERS[chosen_idx] if chosen_idx < len(LETTERS) else "?"
        correct = LETTERS[correct_idx] if correct_idx < len(LETTERS) else "?"
        state.answers.append((state.index + 1, chosen, correct, ok))
        if state.mode == "study":
            await query.edit_message_text(
                answered_question_text(state, question, chosen_idx, correct_idx),
                parse_mode=ParseMode.HTML,
            )
            await send_answer_feedback(context, chat_id, question, chosen_idx, correct_idx, ok)
        else:
            await query.edit_message_reply_markup(reply_markup=None)
        state.index += 1
        if state.index >= len(state.questions):
            await finish_quiz(chat_id, context)
        else:
            await send_current_question(chat_id, None, context)


async def text_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    text = update.message.text if update.message else ""
    match = NUMBER_RE.match(text or "")
    if match:
        await start_quiz(update, random_questions(clamp_count(int(match.group(1)))), "study")
        return
    await send_text(update, "Send a number like 20, or use /categories, /keywords, /search, /random.", main_menu())


async def post_init(app: Application) -> None:
    await app.bot.set_my_commands([
        BotCommand("start", "Show help and quick quiz buttons"),
        BotCommand("random", "Start a random quiz: /random 20 exam"),
        BotCommand("categories", "Browse categories"),
        BotCommand("keywords", "Browse/search keywords: /keywords fever"),
        BotCommand("search", "Search text/topic: /search sepsis 20"),
        BotCommand("stop", "Stop the current quiz"),
    ])


def main() -> None:
    app = Application.builder().token(TELEGRAM_BOT_TOKEN).post_init(post_init).build()
    app.add_handler(CommandHandler(["start", "help"], help_cmd))
    app.add_handler(CommandHandler("random", random_cmd))
    app.add_handler(CommandHandler("search", search_cmd))
    app.add_handler(CommandHandler("categories", categories_cmd))
    app.add_handler(CommandHandler(["keywords", "tags"], keywords_cmd))
    app.add_handler(CommandHandler("stop", stop_cmd))
    app.add_handler(CallbackQueryHandler(callback_handler))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, text_handler))
    app.run_polling(allowed_updates=Update.ALL_TYPES)


if __name__ == "__main__":
    main()