"""Stage 2, clean-up: turn the text layer of a chapter into clean paragraphs.
Every rule is deterministic. Every change and every word the dictionary does not know goes to warnings.log.
Author words are never rewritten, only layout artifacts are removed or repaired."""
import os
import re

from spellchecker import SpellChecker

import bookio
import pageparts
import pdftext
import structure

WORD_RE = re.compile(r"[A-Za-z]+(?:'[A-Za-z]+)?")
_dict = SpellChecker()


def known(word):
    return bool(_dict.known([word.lower()]))


def vocabulary(units):
    """Words and hyphenated words that occur inline (not split by a line break) anywhere in the book."""
    inline = set()
    for u in units:
        for ln in u["text"].split("\n"):
            toks = re.findall(r"[A-Za-z]+(?:-[A-Za-z]+)*", ln)
            if ln.rstrip().endswith("-") and toks:
                toks = toks[:-1]
            inline.update(t.lower() for t in toks)
    return inline


def parts_known(word, inline):
    return all(known(p) or p.lower() in inline for p in word.split("-") if p)


def join_hyphen(left, frag, inline):
    """Decide how to join 'left-' + 'frag'. Returns (joined_word, reason, resolved)."""
    joined, hyph = left + frag, left + "-" + frag
    if hyph.lower() in inline:
        return hyph, "hyphenated form occurs elsewhere in the book", True
    if parts_known(joined, inline) and (known(joined.split("-")[-1]) or joined.lower() in inline):
        return joined, "joined word is a known word", True
    if known(left.split("-")[-1]) and known(frag):
        return hyph, "both fragments are words and the joined form is unknown, hyphen kept", True
    return joined, "neither form is known, joined by default", False


def split_dropcap(lines, i):
    """A lone capital letter followed by a line that starts with the rest of the first word."""
    if re.fullmatch(r"[A-Z]", lines[i]["text"].strip()) and i + 1 < len(lines):
        nxt = lines[i + 1]["text"].strip()
        if re.match(r"[A-Z]{2,}\b", nxt) or re.match(r"[A-Z]{2,}$", nxt):
            return True
    return False


def smallcaps_lead(text):
    """'ALL MY LIFE I have' -> 'All my life I have' (leading run of capitals after a drop cap)."""
    toks = text.split(" ")
    n = 0
    while n < len(toks) and toks[n].strip(".,;:!?\"'").isupper():
        n += 1
    out = []
    for k, t in enumerate(toks):
        if k < n and len(t.strip(".,;:!?\"'")) > 1:
            t = t.lower()
        if k == 0:
            t = t[:1].upper() + t[1:]
        out.append(t)
    return " ".join(out)


def unit_lines(unit, title, heading, first_unit, notes):
    """Body lines of one unit as dicts {text, indent, where}; removed parts are noted."""
    raw = [l.rstrip() for l in unit["text"].split("\n")]
    geom = unit.get("geom") or [None] * len(raw)
    where = "pdf p.%d%s" % (unit["pdf_page"], unit["half"] or "")
    drop = set()
    z = pageparts.header_zone(raw, title, [heading["title"]] if heading else [])
    if z:
        drop.update(range(z[0], z[1]))
        what = "page number %d" % z[2] if z[2] is not None else "no page number"
        notes.append(("header-removed", where, "running header '%s' (%s)" % (" ".join(l.strip() for l in raw[z[0]:z[1]]), what)))
    f = pageparts.footer_number(raw)
    if f:
        drop.add(f[0])
        notes.append(("page-number-removed", where, "footer page number %d" % f[1]))
    if first_unit and heading:
        drop.update(range(0, heading["title_end"]))
        notes.append(("heading-removed", where, "chapter heading 'CHAPTER %d / %s' (spoken header is added in the script)" % (heading["number"], heading["title"])))
    # last line of running text: the lines after it (all capitals only) are picture captions, not headings
    body = [i for i, l in enumerate(raw) if i not in drop and any(c.islower() for c in l)]
    last_body = body[-1] if body else -1
    out = []
    for i, ln in enumerate(raw):
        if i in drop or not ln.strip():
            continue
        if pageparts.CAPTION_RE.match(ln):
            notes.append(("caption-removed", where, "figure caption '%s' (listed in figures/skipped_figures.md)" % ln.strip()))
            continue
        if CAPTION_MIXED_RE.match(ln):
            notes.append(("caption-removed", where, "figure caption '%s' (marker not in capitals)" % ln.strip()))
            continue
        letters = [c for c in ln if c.isalpha()]
        if len(letters) >= 3 and not any(c.islower() for c in letters):
            if i > last_body:
                notes.append(("caption-removed", where, "all-capitals line '%s' below the text (picture caption or label)" % ln.strip()))
                continue
            spoken = bookio.title_case(" ".join(ln.split()))
            notes.append(("heading-kept", where, "all-capitals line above the text kept as a section heading: '%s' -> '%s'" % (ln.strip(), spoken)))
            ln = " " * (len(ln) - len(ln.lstrip())) + spoken
        g = geom[i] if i < len(geom) else None
        out.append({"text": ln.strip(), "indent": len(ln) - len(ln.lstrip()), "where": where,
                    "unit": (unit["pdf_page"], unit["half"]), "x": g["x"] if g else None, "y": g["y"] if g else None})
    return out


INDENT_PT = 4.5    # a paragraph's first line sits at least this many points right of the line before it
GAP_PT = 2.5       # or this many points lower than the page's normal line spacing
SAMELINE_CAP = re.compile(r"^([A-Z])\s{2,}([A-Z]{2,})\b")


def unit_stats(lines):
    """Per page unit: normal line pitch and the local left margin, read from line positions in points."""
    stats = {}
    for u in {l["unit"] for l in lines}:
        mine = [l for l in lines if l["unit"] == u and l["x"] is not None]
        gaps = sorted(b["y"] - a["y"] for a, b in zip(mine, mine[1:]) if 8 < b["y"] - a["y"] < 20)
        stats[u] = {"pitch": gaps[(len(gaps) - 1) // 2] if gaps else 13.0, "lines": mine}
    return stats


def local_margin(stats, ln):
    """Left margin near the first line of a unit: the leftmost of its next few lines (the page may be skewed)."""
    near = [l["x"] for l in stats[ln["unit"]]["lines"] if l is not ln and 0 < l["y"] - ln["y"] < 90]
    return min(near) if near else ln["x"]


def starts_paragraph(prev, ln, stats):
    """True if the line opens a paragraph. Evidence is geometric: indented relative to the line before it, or
    set lower than normal line spacing. A line with no position, or one that begins in lower case, never starts one."""
    if ln["x"] is None or ln["text"][:1].islower():
        return False
    if prev is not None and prev["unit"] == ln["unit"] and prev["x"] is not None:
        return ln["x"] - prev["x"] >= INDENT_PT or ln["y"] - prev["y"] >= stats[ln["unit"]]["pitch"] + GAP_PT
    return ln["x"] - local_margin(stats, ln) >= INDENT_PT


def paragraphs_from(lines, inline, notes):
    """Group lines into paragraphs and repair drop caps. With line positions (the normal case) a paragraph break is
    an indent or a vertical gap; without them (older callers) it is the layout indent of the first line."""
    geometric = any(l.get("x") is not None for l in lines)
    stats = unit_stats(lines) if geometric else None
    base = min(l["indent"] for l in lines)
    i, paras, cur, cap_indent, wrap, prev = 0, [], None, None, False, None
    while i < len(lines):
        ln = dict(lines[i])
        is_cap = split_dropcap(lines, i)
        same = SAMELINE_CAP.match(ln["text"]) if not is_cap else None
        if geometric:
            if wrap and ln["x"] is not None and ln["x"] - local_margin(stats, ln) >= INDENT_PT and not same:
                is_start = False           # still beside the drop cap
            else:
                wrap = False
                is_start = cur is None or starts_paragraph(prev, ln, stats)
        else:
            is_start = ln["indent"] - base >= 1 and not (cap_indent is not None and ln["indent"] == cap_indent)
        if same and (cur is None or is_start):
            letter, rest = same.group(1), ln["text"][same.end(1):].strip()
            first = same.group(2)
            if known(letter + first) or not (letter in "AI" and known(first)):
                merged, how = letter + rest, "joined"
            else:
                merged, how = letter + " " + rest, "kept as its own word"
            notes.append(("dropcap", ln["where"], "drop cap '%s' + '%s' on one line: %s" % (letter, rest[:12], how)))
            ln.update(text=merged, indent=base + 1, smallcaps=True)
            cur = [ln]
            paras.append(cur)
            wrap, prev = True, lines[i]
            i += 1
            continue
        if is_cap:
            letter, frag = ln["text"], lines[i + 1]
            first = re.match(r"[A-Z]+", frag["text"]).group(0)
            if known(letter + first) or not (letter in "AI" and known(first)):
                merged = letter + frag["text"]
                how = "joined"
            else:
                merged = letter + " " + frag["text"]
                how = "kept as its own word"
            notes.append(("dropcap", ln["where"], "drop cap '%s' + '%s': %s" % (letter, frag["text"][:12], how)))
            ln = dict(frag, text=merged, indent=base + 1, where=ln["where"], smallcaps=True)
            cap_indent = frag["indent"]
            i += 2
            cur = [ln]
            paras.append(cur)
            wrap, prev = True, frag
            continue
        if ln["indent"] != cap_indent:
            cap_indent = None
        if is_start or cur is None:
            cur = [ln]
            paras.append(cur)
        else:
            cur.append(ln)
        prev = lines[i]
        i += 1
    return paras


def join_lines(para, inline, notes):
    text = para[0]["text"]
    for ln in para[1:]:
        nxt = ln["text"]
        m = re.search(r"([A-Za-z]+(?:-[A-Za-z]+)*)-$", text)
        f = re.match(r"([a-z]+)(.*)$", nxt)
        if m and f:
            joined, reason, ok = join_hyphen(m.group(1), f.group(1), inline)
            notes.append(("hyphen-joined" if ok else "hyphen-unresolved", ln["where"],
                          "'%s-' + '%s' -> '%s' (%s)" % (m.group(1), f.group(1), joined, reason)))
            text = text[:m.start()] + joined + f.group(2)
        else:
            text = text + " " + nxt
    return re.sub(r"\s+", " ", text).strip()


ELISION = {"Tis", "Twas", "Twere", "Twill", "Twould"}
CAPTION_MIXED_RE = re.compile(r"^\s*(?:Left|Right|Above|Below|Opposite|Facing)\s*:\s*[A-Z][A-Z .,'\-]+\s*$")
MIDLINE_HYPHEN_RE = re.compile(r"\b([A-Za-z]+)- ([a-z]+)\b")
LABEL_RE = re.compile(r"\s\d\.\s+[A-Z]{3,}(?:\s+[A-Z]{3,})*(?=\s+[a-z])")
DIGIT_ONE_RE = re.compile(r"(?<![\d.,$#/-])\b1 (?=[a-z]+\b)")


def strip_inline_labels(text, where, notes):
    """The book sets '5. JOY' style labels into the middle of a sentence (a typeset inset, not words of the sentence).
    Left in, they would be read aloud inside the sentence, so they are removed and logged."""
    def drop(m):
        notes.append(("inline-label-removed", where, "typeset label '%s' sat inside a sentence" % m.group(0).strip()))
        return ""
    return LABEL_RE.sub(drop, text)


def fix_midline_hyphens(text, where, notes, inline):
    """'worth- while' inside a line: the scan kept a line-end hyphen after the line was re-flowed.
    Joined only when the answer is certain (the hyphenated form occurs elsewhere in the book, or the joined
    word is a dictionary word). Otherwise logged and left as it is."""
    def fix(m):
        joined, reason, ok = join_hyphen(m.group(1), m.group(2), inline)
        certain = ok and (reason.startswith("hyphenated form") or reason.startswith("joined word"))
        notes.append(("hyphen-joined" if certain else "hyphen-unresolved", where,
                      "mid-line '%s- %s' -> %s (%s)" % (m.group(1), m.group(2), "'%s'" % joined if certain else "not changed", reason)))
        return joined if certain else m.group(0)
    return MIDLINE_HYPHEN_RE.sub(fix, text)


def merge_orphan_quotes(paras, chapter, notes):
    """[(text, where)] -> same, where a paragraph that opens with a lone quote mark and a lower-case word is
    joined to the paragraph before it. No paragraph starts like that: the page broke the line after an opening
    quoted word, as in: They "milked'  /  " every operation. The stray ' closing the quoted word is made a '"'."""
    out = []
    for text, where in paras:
        if out and re.match(r'"\s+[a-z]', text):
            prev, prev_where = out[-1]
            fixed = prev
            if prev.count('"') % 2 == 1 and re.search(r"[A-Za-z]'$", prev):
                fixed = prev[:-1] + '"'
            out[-1] = (fixed + " " + text[1:].strip(), prev_where)
            notes.append(("paragraph-joined", where, "paragraph began with a lone quote mark and a lower-case word, "
                          "so it continues the paragraph before it ('...%s' + '%s...')" % (prev[-14:], text[:14])))
        else:
            out.append((text, where))
    return out


def fix_quotes(text, where, notes):
    """Quote marks the scan got wrong. Fixed only when the rule is certain; otherwise the paragraph is logged.
      - an unclosed double quote that ends in punctuation + ' : the ' is a misread closing quote.
      - a paragraph that opens with ' + capital letter (not an elision such as 'Tis): a misread opening
        double quote, as at the start of each paragraph of a long speech.
      - a closing double quote (after punctuation) with no opening quote before it in the paragraph is logged only."""
    if text.count('"') % 2 == 1 and re.search(r"[.?!]'$", text):
        notes.append(("quote-fixed", where, "closing quote misread as apostrophe at '...%s'" % text[-12:]))
        text = text[:-1] + '"'
    if re.match(r"\u2022[A-Z]", text):
        notes.append(("quote-fixed", where, "opening quote misread as a bullet at '%s...'" % text[:16]))
        text = '"' + text[1:]
    m = re.match(r"'([A-Z][a-z]*)", text)
    if m and m.group(1) not in ELISION:
        notes.append(("quote-fixed", where, "opening quote misread as apostrophe at '%s...'" % text[:16]))
        text = '"' + text[1:]
    opened = False
    for q in re.finditer('"', text):
        before = text[q.start() - 1] if q.start() else ""
        if before in ("", " ", "(", "\u2014"):
            opened = True
        elif before == "'":
            opened = not opened
        elif opened:
            opened = False
        else:
            notes.append(("quote-unmatched", where, "closing quote with no opening quote before it: '...%s' (not changed)" % text[max(0, q.start() - 14):q.start() + 1]))
    for m in DIGIT_ONE_RE.finditer(text):
        notes.append(("digit-one", where, "'1' before a lowercase word, probably the letter I misread: '...%s' (not changed)" % text[max(0, m.start() - 8):m.end() + 8]))
    return text


def odd_words(text, where, notes):
    for w in WORD_RE.findall(text):
        core = re.sub(r"'s$", "", w)
        if len(core) > 1 and not known(core):
            notes.append(("odd-word", where, "'%s' is not in the dictionary%s" % (w, " (capitalised: probably a name)" if w[0].isupper() else "")))


def run(book_dir, book, pdf, chapter_numbers):
    """Writes text/chapter-NN.txt for the chosen chapters. Returns warnings."""
    units = pdftext.units(pdf)
    heads = {h["number"]: h for h in structure.find_headings(units)}
    inline = vocabulary(units)
    os.makedirs(os.path.join(book_dir, "text"), exist_ok=True)
    notes = []
    for ch in book["chapters"]:
        if ch["number"] not in chapter_numbers:
            continue
        lines = []
        for ui in range(ch["first_unit"], ch["last_unit"] + 1):
            lines += unit_lines(units[ui], book["title"], heads[ch["number"]], ui == ch["first_unit"], notes)
        paras = paragraphs_from(lines, inline, notes)
        joined = []
        for p in paras:
            t = join_lines(p, inline, notes)
            if p[0].get("smallcaps"):
                before, t = t, smallcaps_lead(t)
                if before != t:
                    notes.append(("smallcaps", p[0]["where"], "'%s...' -> '%s...'" % (before[:14], t[:14])))
            joined.append((t, p[0]["where"]))
        joined = merge_orphan_quotes(joined, ch["number"], notes)
        texts = []
        for n, (t, _) in enumerate(joined, start=1):
            where = "ch%d para %d" % (ch["number"], n)
            t = strip_inline_labels(t, where, notes)
            t = fix_midline_hyphens(t, where, notes, inline)
            t = fix_quotes(t, where, notes)
            odd_words(t, where, notes)
            texts.append(t)
        name = "chapter-%02d.txt" % ch["number"]
        with open(os.path.join(book_dir, "text", name), "w", encoding="utf-8") as f:
            f.write("\n\n".join(texts) + "\n")
        ch["text_file"] = "text/" + name
        ch["words"] = sum(len(t.split()) for t in texts)
        ch["paragraphs"] = len(texts)
    return notes