#!/usr/bin/env python3
# prep.py
#
# End-to-end, fast preprocessing:
# 1) Read markdown from scraped.db (status_code=200, non-empty).
# 2) Clean + annotate anchors:
#       [anchor](url)   -> [LINK_START]anchor[LINK_END]
#       [anchor][ref]   -> [LINK_START]anchor[LINK_END]
#       ![alt](url)     -> removed
#       <http...>       -> removed
#       bare URLs       -> removed
#    Output: train_clean.csv (no header, one fully-quoted line per doc, no newlines).
# 3) Tokenize with microsoft/mdeberta-v3-base; align labels from LINK spans.
# 4) Split docs 9:1 (doc-level), windowize to max_length with doc_stride.
#    Output: train_windows.jsonl, val_windows.jsonl, prep_summary.txt
#
# Usage:
#   python prep.py --db scraped.db --max_length 512 --doc_stride 128 --val_ratio 0.1 --seed 42 --workers auto

import os
import re
import csv
import sys
import json
import math
import sqlite3
import random
import argparse
from typing import List, Tuple, Dict
from concurrent.futures import ProcessPoolExecutor

import numpy as np
from tqdm import tqdm

try:
    from transformers import AutoTokenizer
except ImportError:
    print("ERROR: transformers not installed. pip install transformers", file=sys.stderr); sys.exit(1)

LINK_START = "[LINK_START]"
LINK_END = "[LINK_END]"

# ---------- Fast regex (no nested bracket support) ----------
IMG_INLINE_RE   = re.compile(r'!\[[^\]]*\]\([^)]*\)')
INLINE_LINK_RE  = re.compile(r'\[([^\]]+)\]\([^)]*\)')       # [text](url) -> capture text
REF_LINK_RE     = re.compile(r'\[([^\]]+)\]\[[^\]]+\]')      # [text][id] -> capture text
REF_DEF_RE      = re.compile(r'^[ \t]{0,3}\[[^\]]+\]:\s+\S+.*$', re.MULTILINE)  # definition lines
AUTOLINK_RE     = re.compile(r'<https?[^>]+>')
BARE_URL_RE     = re.compile(r'https?://\S+|www\.\S+')
CODE_TICKS_RE   = re.compile(r'`+')
EMPH_RE         = re.compile(r'[*]+')  # FIX: do NOT strip underscores (protects [LINK_START]/[LINK_END])
HEAD_RE         = re.compile(r'^[ \t]*#+[ \t]*', re.MULTILINE)
QUOTE_RE        = re.compile(r'^(>+\s*)+', re.MULTILINE)
WS_RE           = re.compile(r'\s+')

def annotate_one(md_text: str) -> Tuple[str, int]:
    """Return (single-line annotated text, has_marker[0/1])."""
    if not md_text:
        return "", 0
    t = md_text

    # Order matters
    t = IMG_INLINE_RE.sub('', t)
    t = INLINE_LINK_RE.sub(lambda m: f"{LINK_START}{m.group(1)}{LINK_END}", t)
    t = REF_LINK_RE.sub(lambda m: f"{LINK_START}{m.group(1)}{LINK_END}", t)
    t = REF_DEF_RE.sub('', t)
    t = AUTOLINK_RE.sub('', t)
    t = BARE_URL_RE.sub('', t)

    # Lightweight markdown cleanup; keep brackets and underscores (markers)
    t = CODE_TICKS_RE.sub('', t)
    t = EMPH_RE.sub('', t)
    t = HEAD_RE.sub('', t)
    t = QUOTE_RE.sub('', t)

    # Collapse whitespace and ensure single line
    t = WS_RE.sub(' ', t).strip()

    has = 1 if (LINK_START in t and LINK_END in t) else 0
    return t, has

# ---------- Spans & windowing ----------
def strip_and_get_spans(s: str) -> Tuple[str, List[Tuple[int, int]]]:
    """Remove LINK markers and return (plain_text, spans) in char offsets."""
    spans: List[Tuple[int, int]] = []
    out: List[str] = []
    i = 0
    n = len(s)
    in_link = False
    start_pos = -1
    while i < n:
        if s.startswith(LINK_START, i):
            if not in_link:
                in_link = True
                start_pos = len(out)
            i += len(LINK_START); continue
        if s.startswith(LINK_END, i):
            if in_link:
                in_link = False
                end_pos = len(out)
                if end_pos > start_pos >= 0:
                    spans.append((start_pos, end_pos))
                start_pos = -1
            i += len(LINK_END); continue
        out.append(s[i]); i += 1
    return "".join(out), spans

def labels_from_spans(offset_mapping: List[Tuple[int, int]], spans: List[Tuple[int, int]]) -> List[int]:
    """Binary label 1 if token overlaps any span by >=1 char, else 0."""
    labels: List[int] = []
    spans = sorted(spans)
    for ts, te in offset_mapping:
        if ts == te:
            labels.append(0); continue
        lab = 0
        for ss, se in spans:
            if te <= ss: break
            if ts >= se: continue
            lab = 1; break
        labels.append(lab)
    return labels

def windowize_ids_and_labels(
    input_ids_no_special: List[int],
    labels_no_special: List[int],
    tokenizer: AutoTokenizer,
    max_length: int,
    doc_stride: int
) -> Tuple[List[List[int]], List[List[int]], List[List[int]]]:
    """Slice long sequences to windows with specials (<= max_length)."""
    assert len(input_ids_no_special) == len(labels_no_special)
    specials = tokenizer.num_special_tokens_to_add(pair=False)
    cap = max_length - specials
    if cap <= 0:
        raise ValueError(f"max_length too small; specials={specials}")

    def pack(ids_no_sp: List[int], labs_no_sp: List[int]):
        ids_with = tokenizer.build_inputs_with_special_tokens(ids_no_sp)
        attn = [1] * len(ids_with)
        if specials == 2:
            labs_with = [0] + labs_no_sp + [0]
        else:
            pad_n = len(ids_with) - len(labs_no_sp)
            labs_with = [0] * pad_n
            if pad_n >= 1:
                labs_with = [0] + labs_no_sp + [0] * (pad_n - 1)
            else:
                labs_with = labs_no_sp[:len(ids_with)]
        return ids_with[:max_length], attn[:max_length], labs_with[:max_length]

    if len(input_ids_no_special) <= cap:
        ids_w, attn_w, labs_w = pack(input_ids_no_special, labels_no_special)
        return [ids_w], [attn_w], [labs_w]

    step = max(cap - doc_stride, 1)
    out_ids: List[List[int]] = []
    out_attn: List[List[int]] = []
    out_labs: List[List[int]] = []
    start = 0
    total = len(input_ids_no_special)
    while start < total:
        end = min(start + cap, total)
        ids_slice = input_ids_no_special[start:end]
        labs_slice = labels_no_special[start:end]
        ids_w, attn_w, labs_w = pack(ids_slice, labs_slice)
        out_ids.append(ids_w); out_attn.append(attn_w); out_labs.append(labs_w)
        if end == total: break
        start += step
    return out_ids, out_attn, out_labs

# ---------- DB read ----------
def read_markdown_from_db(db_path: str) -> List[str]:
    conn = sqlite3.connect(db_path)
    try:
        cur = conn.cursor()
        cur.execute("""
            SELECT full_markdown_content
            FROM scraped_data
            WHERE status_code = 200
              AND full_markdown_content IS NOT NULL
              AND TRIM(full_markdown_content) != ''
        """)
        rows = cur.fetchall()
        return [r[0] if isinstance(r[0], str) else str(r[0]) for r in rows]
    finally:
        conn.close()

# ---------- Main workflow ----------
def main():
    p = argparse.ArgumentParser(description="Fast end-to-end preprocessing for link token classification.")
    p.add_argument("--db", default="scraped.db", help="SQLite DB path (table scraped_data).")
    p.add_argument("--output_csv", default="train_clean.csv", help="Output cleaned CSV (quoted, one line/doc).")
    p.add_argument("--tokenizer", default="microsoft/mdeberta-v3-base", help="HF tokenizer.")
    p.add_argument("--max_length", type=int, default=512, help="Max tokens incl specials.")
    p.add_argument("--doc_stride", type=int, default=128, help="Overlap on content tokens.")
    p.add_argument("--val_ratio", type=float, default=0.1, help="Validation ratio by document.")
    p.add_argument("--seed", type=int, default=42, help="Random seed for split.")
    p.add_argument("--batch_size", type=int, default=64, help="Tokenization batch size.")
    p.add_argument("--workers", default="auto", help="Annotation worker count: int or 'auto'.")
    args = p.parse_args()

    script_dir = os.path.dirname(os.path.abspath(__file__))
    db_path = os.path.join(script_dir, args.db)
    out_csv = os.path.join(script_dir, args.output_csv)
    if not os.path.isfile(db_path):
        print(f"ERROR: DB not found: {db_path}", file=sys.stderr); sys.exit(1)

    # 1) Read markdown
    print(f"[1/4] Read from DB: {args.db}")
    md_rows = read_markdown_from_db(db_path)
    n_docs = len(md_rows)
    print(f"     Rows: {n_docs}")

    # 2) Clean + annotate (multiprocessing)
    print(f"[2/4] Clean + annotate -> {args.output_csv}")
    workers = os.cpu_count() if args.workers == "auto" else int(args.workers)
    markers = 0
    written = 0

    with open(out_csv, "w", encoding="utf-8", newline="") as f_out:
        writer = csv.writer(f_out, quoting=csv.QUOTE_ALL)
        with ProcessPoolExecutor(max_workers=workers) as ex:
            for txt, has in tqdm(ex.map(annotate_one, md_rows, chunksize=512), total=n_docs, unit="doc", desc="Annotating"):
                if not txt:
                    continue
                if '\n' in txt or '\r' in txt:
                    txt = WS_RE.sub(' ', txt).strip()
                writer.writerow([txt])
                written += 1
                markers += has

    if written == 0:
        print("ERROR: No documents written after cleaning.", file=sys.stderr); sys.exit(1)
    print(f"     Written: {written}  |  With LINK markers: {markers}")

    # 3) Tokenize + labels + split + windowize
    print(f"[3/4] Tokenize + align + split + windowize (tokenizer={args.tokenizer})")
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True)
    specials = tokenizer.num_special_tokens_to_add(pair=False)
    cap = args.max_length - specials
    if cap <= 0:
        print(f"ERROR: max_length too small for specials={specials}", file=sys.stderr); sys.exit(1)

    texts = []
    with open(out_csv, "r", encoding="utf-8", newline="") as f_in:
        rdr = csv.reader(f_in, quoting=csv.QUOTE_ALL)
        for row in rdr:
            texts.append(row[0])
    num_docs = len(texts)

    plain_texts: List[str] = []
    spans_all: List[List[Tuple[int, int]]] = []
    for t in tqdm(texts, total=num_docs, unit="doc", desc="Extract spans"):
        plain, spans = strip_and_get_spans(t)
        plain_texts.append(plain)
        spans_all.append(spans)

    input_ids_no_sp: List[List[int]] = []
    offsets_all: List[List[Tuple[int, int]]] = []
    for i in tqdm(range(0, num_docs, args.batch_size), unit="batch", desc="Tokenize"):
        batch = plain_texts[i:i+args.batch_size]
        enc = tokenizer(
            batch,
            add_special_tokens=False,
            return_offsets_mapping=True,
            return_attention_mask=False,
            return_token_type_ids=False,
            truncation=False,
        )
        input_ids_no_sp.extend(enc["input_ids"])
        offsets_all.extend([[(int(a), int(b)) for (a, b) in off] for off in enc["offset_mapping"]])

    labels_no_sp: List[List[int]] = []
    total_tokens = 0
    pos_tokens = 0
    for offs, spans in tqdm(zip(offsets_all, spans_all), total=num_docs, unit="doc", desc="Align labels"):
        labs = labels_from_spans(offs, spans)
        labels_no_sp.append(labs)
        total_tokens += len(labs)
        if labs:
            pos_tokens += int(np.sum(labs))

    idx = list(range(num_docs))
    random.Random(args.seed).shuffle(idx)
    val_n = max(1, int(round(num_docs * args.val_ratio)))
    val_set = set(idx[:val_n])

    train_out_path = os.path.join(script_dir, "train_windows.jsonl")
    val_out_path = os.path.join(script_dir, "val_windows.jsonl")
    train_out = open(train_out_path, "w", encoding="utf-8")
    val_out = open(val_out_path, "w", encoding="utf-8")

    train_windows = 0
    val_windows = 0
    train_win_with_link = 0
    val_win_with_link = 0
    exceeding_docs = 0

    for doc_id in tqdm(range(num_docs), unit="doc", desc="Windowize+write"):
        ids = input_ids_no_sp[doc_id]
        labs = labels_no_sp[doc_id]
        if len(ids) + specials > args.max_length:
            exceeding_docs += 1
        ids_ws, attn_ws, labs_ws = windowize_ids_and_labels(ids, labs, tokenizer, args.max_length, args.doc_stride)
        target = val_out if doc_id in val_set else train_out
        for w_id, (iw, aw, lw) in enumerate(zip(ids_ws, attn_ws, labs_ws)):
            if any(x == 1 for x in lw):
                if doc_id in val_set: val_win_with_link += 1
                else: train_win_with_link += 1
            rec = {"doc_id": int(doc_id), "window_id": int(w_id), "input_ids": iw, "attention_mask": aw, "labels": lw}
            target.write(json.dumps(rec, ensure_ascii=False) + "\n")
        if doc_id in val_set: val_windows += len(ids_ws)
        else: train_windows += len(ids_ws)

    train_out.close(); val_out.close()

    # 4) Summary
    pos_rate = (pos_tokens / total_tokens) if total_tokens else 0.0
    summary_lines = [
        "=== prep.py Summary ===",
        f"DB: {args.db}",
        f"Output CSV: {args.output_csv}",
        f"Tokenizer: {args.tokenizer}",
        f"max_length: {args.max_length} (specials={specials}, content_capacity={cap})",
        f"doc_stride: {args.doc_stride}",
        f"Documents cleaned: {num_docs}",
        f"Documents exceeding max_length (incl specials): {exceeding_docs}",
        f"Tokens total (no specials): {total_tokens}",
        f"Positive tokens: {pos_tokens} ({pos_rate:.4%})",
        f"Train windows: {train_windows} (with_link={train_win_with_link})",
        f"Val windows:   {val_windows} (with_link={val_win_with_link})",
        f"Train JSONL: train_windows.jsonl",
        f"Val JSONL:   val_windows.jsonl",
    ]
    with open(os.path.join(script_dir, "prep_summary.txt"), "w", encoding="utf-8") as f:
        f.write("\n".join(summary_lines) + "\n")

    print(f"[4/4] Summary -> prep_summary.txt\n" + "\n".join(summary_lines))
    print("Done.")

if __name__ == "__main__":
    main()