from __future__ import annotations

import re
from collections import defaultdict
from difflib import SequenceMatcher
from typing import Any

from .render import render_block
from .source import SourcePackage, TsvWord


TOKEN_RE = re.compile(r"[\w]+(?:['’][\w]+)?", re.UNICODE)


def _tokens(value: str) -> list[str]:
    return [token.casefold().replace("’", "'") for token in TOKEN_RE.findall(value)]


def _bbox(words: list[TsvWord]) -> dict[str, int] | None:
    if not words:
        return None
    left = min(word.left for word in words)
    top = min(word.top for word in words)
    right = max(word.left + word.width for word in words)
    bottom = max(word.top + word.height for word in words)
    return {"x": left, "y": top, "width": right - left, "height": bottom - top}


def align_page_blocks(page: dict[str, Any], source: SourcePackage) -> dict[str, Any]:
    """Best-effort exact token alignment; failure never blocks transcription."""
    page_number = page["source"]["pdf_page"]
    ocr_words = source.tsv_words(page_number)
    ocr_tokens = [token for word in ocr_words for token in _tokens(word.text)]
    ocr_token_words = [word for word in ocr_words for _ in _tokens(word.text)]

    transcription_tokens: list[str] = []
    token_blocks: list[str] = []
    totals: dict[str, int] = {}
    for block in page["blocks"]:
        tokens = _tokens(render_block(block, diplomatic=False))
        totals[block["id"]] = len(tokens)
        transcription_tokens.extend(tokens)
        token_blocks.extend([block["id"]] * len(tokens))

    matcher = SequenceMatcher(
        a=transcription_tokens,
        b=ocr_tokens,
        autojunk=False,
    )
    matched_words: dict[str, list[TsvWord]] = defaultdict(list)
    matched_indices: dict[str, set[int]] = defaultdict(set)
    for match in matcher.get_matching_blocks():
        for offset in range(match.size):
            transcription_index = match.a + offset
            ocr_index = match.b + offset
            block_id = token_blocks[transcription_index]
            matched_words[block_id].append(ocr_token_words[ocr_index])
            matched_indices[block_id].add(transcription_index)

    for block in page["blocks"]:
        block_id = block["id"]
        total = totals[block_id]
        words = matched_words[block_id]
        unique_words = list(dict.fromkeys(words))
        matched = len(matched_indices[block_id])
        if total == 0:
            status = "unaligned"
        elif matched == total:
            status = "aligned"
        elif matched * 2 >= total:
            status = "partial"
        else:
            status = "unaligned"
        block["evidence"] = {
            "status": status,
            "method": "exact_normalized_token_sequence" if total else None,
            "ocr_word_ids": [word.word_id for word in unique_words],
            "bbox": _bbox(unique_words),
            "matched_tokens": matched,
            "total_tokens": total,
        }
    return page

