"""An exam's published shape, and drawing a paper that matches it.

A real paper is not a uniform draw from a bank. The ABP publishes that 12% of a
general paediatrics exam is preventive care and 2% is rheumatology; a block
drawn at random is forty coin flips, and tells a learner nothing about how they
would do on the day.

The blueprint is a mapping, not a reshaping. Our taxonomy is arranged for
studying and a board's outline is arranged for examining, and both are right —
so a blueprint line names its categories rather than the tree being bent to fit.
"""
from decimal import Decimal

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.models.exam import BlueprintCategoryLink, ExamBlueprint, QuestionExamLink
from app.models.question import Question
from app.models.question_category import QuestionCategory, QuestionCategoryLink


def _descendants(db: Session, category_ids: list[int]) -> set[int]:
    """Every category at or below the given ones."""
    found = set(category_ids)
    frontier = list(category_ids)
    while frontier:
        rows = db.query(QuestionCategory.id).filter(QuestionCategory.parent_id.in_(frontier)).all()
        nxt = [row[0] for row in rows if row[0] not in found]
        found.update(nxt)
        frontier = nxt
    return found


def domains(db: Session, exam_id: int) -> list[ExamBlueprint]:
    """The weighted lines — the domains, not the subdomains beneath them."""
    return (db.query(ExamBlueprint)
            .filter(ExamBlueprint.exam_id == exam_id, ExamBlueprint.parent_id.is_(None))
            .order_by(ExamBlueprint.sort_order, ExamBlueprint.id).all())


def categories_for(db: Session, blueprint_id: int) -> set[int]:
    direct = [row[0] for row in db.query(BlueprintCategoryLink.category_id)
              .filter(BlueprintCategoryLink.blueprint_id == blueprint_id).all()]
    return _descendants(db, direct) if direct else set()


def question_ids_for(db: Session, exam_id: int, category_ids: set[int],
                     predicate=None) -> list[int]:
    """Serveable questions of this exam that fall in these categories."""
    if not category_ids:
        return []
    in_exam = select(QuestionExamLink.question_id).where(QuestionExamLink.exam_id == exam_id)
    linked = select(QuestionCategoryLink.question_id).where(
        QuestionCategoryLink.category_id.in_(category_ids))
    query = db.query(Question.id).filter(
        Question.id.in_(in_exam),
        Question.question_category_id.in_(category_ids) | Question.id.in_(linked),
    )
    if predicate is not None:
        query = query.filter(predicate)
    return [row[0] for row in query.all()]


def allocate(weights: dict[int, Decimal], total: int, available: dict[int, int]) -> dict[int, int]:
    """How many questions each domain should contribute.

    Largest-remainder rather than rounding each share independently: rounding
    twenty-four percentages to whole questions loses or gains several, and a
    forty-question block that returns thirty-six is a bug the learner sees.

    A domain that cannot supply its share gives the shortfall back, and it is
    redistributed over the domains that can — so a thin corner of the bank
    shrinks the paper's accuracy rather than its length.
    """
    if total <= 0 or not weights:
        return {}
    share = sum(weights.values())
    if share <= 0:
        return {}

    exact = {key: (Decimal(total) * weight / share) for key, weight in weights.items()}
    plan = {key: int(value) for key, value in exact.items()}
    # Largest remainder takes the seats whole numbers left over.
    remainder = total - sum(plan.values())
    for key, _ in sorted(exact.items(), key=lambda kv: kv[1] - int(kv[1]), reverse=True):
        if remainder <= 0:
            break
        plan[key] += 1
        remainder -= 1

    # Give back what cannot be supplied, then re-offer it to whoever has room.
    spare = 0
    for key in list(plan):
        cap = available.get(key, 0)
        if plan[key] > cap:
            spare += plan[key] - cap
            plan[key] = cap
    while spare > 0:
        room = [key for key in plan if available.get(key, 0) > plan[key]]
        if not room:
            break
        # In weight order, so the overflow lands where the exam is heaviest.
        room.sort(key=lambda key: weights[key], reverse=True)
        for key in room:
            if spare <= 0:
                break
            plan[key] += 1
            spare -= 1
    return {key: count for key, count in plan.items() if count > 0}


def sample(db: Session, exam_id: int, total: int, predicate=None,
           rng=None) -> tuple[list[int], list[dict]]:
    """A paper shaped like the real one, and the working that produced it.

    Returns the question ids and a per-domain account — what each domain was
    owed and what it could give — so the shortfall is visible rather than the
    block quietly being the wrong shape.
    """
    import random

    rng = rng or random
    lines = domains(db, exam_id)
    if not lines:
        return [], []

    pools: dict[int, list[int]] = {}
    weights: dict[int, Decimal] = {}
    for line in lines:
        if line.weight is None:
            continue
        weights[line.id] = Decimal(line.weight)
        pools[line.id] = question_ids_for(db, exam_id, categories_for(db, line.id), predicate)

    plan = allocate(weights, total, {key: len(ids) for key, ids in pools.items()})

    chosen: list[int] = []
    report = []
    by_id = {line.id: line for line in lines}
    for line_id, weight in sorted(weights.items(), key=lambda kv: kv[1], reverse=True):
        want = plan.get(line_id, 0)
        pool = pools.get(line_id, [])
        picked = rng.sample(pool, min(want, len(pool))) if pool and want else []
        chosen.extend(picked)
        line = by_id[line_id]
        report.append({
            "code": line.code, "title": line.title, "weight": float(weight),
            "asked_for": want, "given": len(picked), "pool": len(pool),
        })
    rng.shuffle(chosen)
    return chosen, report


def coverage(db: Session, exam_id: int) -> list[dict]:
    """What each domain is owed and what the bank can actually supply.

    The number worth knowing before building anything: a domain weighted at 12%
    with nine questions behind it will never carry a paper, and an educator
    should be told that rather than discovering it in a thin block.
    """
    rows = []
    for line in domains(db, exam_id):
        cats = categories_for(db, line.id)
        rows.append({
            "id": line.id,
            "code": line.code,
            "title": line.title,
            "weight": float(line.weight) if line.weight is not None else None,
            "category_ids": sorted(cats),
            "questions": len(question_ids_for(db, exam_id, cats)),
        })
    return rows


def total_weight(db: Session, exam_id: int) -> float:
    value = db.query(func.coalesce(func.sum(ExamBlueprint.weight), 0)).filter(
        ExamBlueprint.exam_id == exam_id, ExamBlueprint.parent_id.is_(None)).scalar()
    return float(value or 0)


def sample_papers(db: Session, exam_id: int, size: int, count: int,
                  predicate=None, rng=None) -> list[tuple[list[int], list[dict]]]:
    """Several papers from one bank, no question sitting in two of them.

    A plan of twelve blocks that each drew independently would ask the same
    preventive-care question four times, because that is the domain with the
    largest share and the smallest pool relative to it. Drawing once and
    dealing out fixes that — but only with a record of what has already gone,
    because two domains can map to the same categories. The board splits
    genitourinary from nephrology and we do not, so a question sits in both
    pools and would otherwise be dealt twice.
    """
    import random

    rng = rng or random
    lines = [line for line in domains(db, exam_id) if line.weight is not None]
    if not lines or size <= 0 or count <= 0:
        return []

    pools: dict[int, list[int]] = {}
    weights: dict[int, Decimal] = {}
    for line in lines:
        ids = question_ids_for(db, exam_id, categories_for(db, line.id), predicate)
        rng.shuffle(ids)
        pools[line.id] = ids
        weights[line.id] = Decimal(line.weight)

    by_id = {line.id: line for line in lines}
    dealt: set[int] = set()

    def take(line_id: int, want: int) -> list[int]:
        out = []
        pool = pools[line_id]
        while pool and len(out) < want:
            candidate = pool.pop()
            if candidate not in dealt:
                dealt.add(candidate)
                out.append(candidate)
        return out

    papers = []
    for _ in range(count):
        # Recounted each round: what a domain can still supply shrinks as it is
        # dealt, and a plan should shorten gracefully rather than repeat.
        room = {key: sum(1 for q in ids if q not in dealt) for key, ids in pools.items()}
        plan = allocate(weights, size, room)
        if not plan:
            break
        chosen: list[int] = []
        report = []
        for line_id, want in plan.items():
            taken = take(line_id, want)
            chosen.extend(taken)
            line = by_id[line_id]
            report.append({"code": line.code, "title": line.title,
                           "weight": float(weights[line_id]),
                           "asked_for": want, "given": len(taken)})
        if not chosen:
            break
        rng.shuffle(chosen)
        report.sort(key=lambda row: row["weight"], reverse=True)
        papers.append((chosen, report))
    return papers


def domain_blocks(db: Session, exam_id: int, size: int, blocks: int,
                  predicate=None, rng=None) -> list[tuple[str, str, list[int]]]:
    """One or more blocks per domain, sized by the weight the board gives it.

    The other shape. A paper is for rehearsal; this is for working through a
    subject, which is what a plan is usually for — so the blocks carry the
    board's own titles in its own order.

    Sized by weight rather than by what happens to be in the bank: chunking
    every question into blocks of forty gave preventive care six blocks and
    the plan a hundred and sixty, which is not a plan. Every domain gets at
    least one block, so nothing the board examines is left out of a plan built
    from its outline.
    """
    import random

    rng = rng or random
    lines = [line for line in domains(db, exam_id) if line.weight is not None]
    if not lines or size <= 0 or blocks <= 0:
        return []

    pools: dict[int, list[int]] = {}
    weights: dict[int, Decimal] = {}
    for line in lines:
        ids = question_ids_for(db, exam_id, categories_for(db, line.id), predicate)
        if not ids:
            continue
        rng.shuffle(ids)
        pools[line.id] = ids
        weights[line.id] = Decimal(line.weight)
    if not pools:
        return []

    # How many blocks each domain is owed, then at least one each.
    share = allocate(weights, blocks, {key: len(ids) for key, ids in pools.items()})
    for key in pools:
        share.setdefault(key, 0)
        share[key] = max(share[key], 1)

    by_id = {line.id: line for line in lines}
    dealt: set[int] = set()
    out = []
    for line in lines:
        if line.id not in pools:
            continue
        want = share[line.id]
        pool = [q for q in pools[line.id] if q not in dealt]
        parts = min(want, max(1, -(-len(pool) // size)))
        for index in range(parts):
            chunk = [q for q in pool[index * size:(index + 1) * size] if q not in dealt]
            if not chunk:
                continue
            dealt.update(chunk)
            suffix = f" ({index + 1} of {parts})" if parts > 1 else ""
            out.append((line.code, f"{by_id[line.id].title}{suffix}", chunk))
    return out