src/services/dataset_prep.py

"""Dataset upload → parse → validate → clean → chat-format conversion.
Runs entirely on CPU (spec P7); PDF/DOCX treated as untrusted input with
config-driven limits (spec §4.2). Token counts are fast estimates (chars/4)
unless a cached tokenizer is available."""

import csv
import hashlib
import io
import json
import pathlib

from src.config_loader import get_configs


class DatasetError(Exception):
    pass


def _limits():
    return get_configs().limits.get("upload", {})


def _check_size(path: pathlib.Path):
    mb = path.stat().st_size / 1e6
    if mb > _limits().get("max_file_mb", 50):
        raise DatasetError(f"File is {mb:.0f} MB — limit is {_limits().get('max_file_mb')} MB.")


def _extract_pdf(path):
    from pypdf import PdfReader  # lazy (spec §12)
    try:
        reader = PdfReader(str(path))
        if reader.is_encrypted:
            raise DatasetError("Encrypted PDFs are not accepted.")
        cap = _limits().get("max_extracted_chars", 5_000_000)
        out, total = [], 0
        for page in reader.pages:
            t = page.extract_text() or ""
            total += len(t)
            if total > cap:
                raise DatasetError(f"Extracted text exceeds {cap} characters.")
            out.append(t)
        return "\n\n".join(out)
    except DatasetError:
        raise
    except Exception as e:  # noqa: BLE001
        raise DatasetError(f"PDF could not be parsed safely: {type(e).__name__}") from e


def _extract_docx(path):
    import docx  # lazy
    try:
        d = docx.Document(str(path))
        text = "\n".join(p.text for p in d.paragraphs)
        if len(text) > _limits().get("max_extracted_chars", 5_000_000):
            raise DatasetError("Extracted text exceeds the configured limit.")
        return text
    except DatasetError:
        raise
    except Exception as e:  # noqa: BLE001
        raise DatasetError(f"DOCX could not be parsed safely: {type(e).__name__}") from e


def _rows_from_structured(path: pathlib.Path):
    suffix = path.suffix.lower()
    text = path.read_text(errors="replace")
    if suffix == ".csv":
        return list(csv.DictReader(io.StringIO(text)))
    if suffix == ".jsonl":
        return [json.loads(l) for l in text.splitlines() if l.strip()]
    if suffix == ".json":
        data = json.loads(text)
        if isinstance(data, dict):
            data = data.get("data") or data.get("rows") or [data]
        return data
    raise DatasetError(f"Unsupported structured format {suffix}")


FIELD_GUESSES = {
    "instruction": ["instruction", "question", "prompt", "input_text", "query"],
    "input": ["input", "context", "passage"],
    "output": ["output", "answer", "response", "completion", "target", "label"],
}


def _guess_fields(row: dict):
    keys = {k.lower(): k for k in row}
    got = {}
    for role, cands in FIELD_GUESSES.items():
        for c in cands:
            if c in keys:
                got[role] = keys[c]
                break
    return got


def _chunk_text(text, target=1500):
    paras = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 60]
    chunks, buf = [], ""
    for p in paras:
        if len(buf) + len(p) > target and buf:
            chunks.append(buf.strip())
            buf = p
        else:
            buf += "\n\n" + p
    if len(buf.strip()) > 200:
        chunks.append(buf.strip())
    return chunks


def prepare(file_path: str, system_prompt: str = "") -> tuple[list[dict], dict]:
    """Returns (records, summary). records = [{"messages": [...]}, ...]"""
    path = pathlib.Path(file_path)
    _check_size(path)
    suffix = path.suffix.lower()
    lim = _limits()

    if suffix in (".csv", ".json", ".jsonl"):
        rows = _rows_from_structured(path)
        if not rows:
            raise DatasetError("No rows found in the file.")
        fields = _guess_fields(rows[0])
        if "output" not in fields or ("instruction" not in fields and "input" not in fields):
            raise DatasetError(
                f"Could not identify instruction/output columns. Found: {list(rows[0].keys())}. "
                f"Rename columns to one of {FIELD_GUESSES['instruction']} + {FIELD_GUESSES['output']}."
            )
        records, has_refs = [], True
        for r in rows:
            user = str(r.get(fields.get("instruction", ""), "")).strip()
            ctx = str(r.get(fields.get("input", ""), "")).strip() if "input" in fields else ""
            out = str(r.get(fields["output"], "")).strip()
            if not (user or ctx) or not out:
                continue
            content = f"{user}\n\n{ctx}".strip()
            msgs = ([{"role": "system", "content": system_prompt}] if system_prompt else [])
            msgs += [{"role": "user", "content": content}, {"role": "assistant", "content": out}]
            records.append({"messages": msgs})
    elif suffix in (".txt", ".pdf", ".docx"):
        text = (path.read_text(errors="replace") if suffix == ".txt"
                else _extract_pdf(path) if suffix == ".pdf" else _extract_docx(path))
        chunks = _chunk_text(text)
        if not chunks:
            raise DatasetError("No usable text extracted.")
        records = [{"messages": (
            [{"role": "system", "content": system_prompt}] if system_prompt else []) + [
            {"role": "user", "content": "Continue writing in the style and subject of this document excerpt:\n\n"
             + c[: len(c) // 2]},
            {"role": "assistant", "content": c[len(c) // 2:]},
        ]} for c in chunks]
        has_refs = False
    else:
        raise DatasetError(f"Unsupported file type {suffix}. Accepted: CSV, JSON, JSONL, TXT, PDF, DOCX.")

    # clean: dedupe + empty filter
    seen, cleaned, dupes = set(), [], 0
    for r in records:
        key = hashlib.sha1(json.dumps(r, sort_keys=True).encode()).hexdigest()
        if key in seen:
            dupes += 1
            continue
        seen.add(key)
        cleaned.append(r)

    if len(cleaned) > lim.get("max_samples", 100_000):
        raise DatasetError(f"{len(cleaned)} samples exceed the limit {lim.get('max_samples')}.")

    lengths = [sum(len(m["content"]) for m in r["messages"]) for r in cleaned]
    est_tokens = int(sum(lengths) / 4)
    if est_tokens > lim.get("max_total_tokens", 20_000_000):
        raise DatasetError(f"Estimated {est_tokens} tokens exceed the limit.")

    summary = {
        "samples": len(cleaned),
        "duplicates_removed": dupes,
        "est_tokens": est_tokens,
        "avg_tokens_per_sample": round(est_tokens / max(len(cleaned), 1), 1),
        "avg_chars": round(sum(lengths) / max(len(lengths), 1), 1),
        "has_reference_answers": has_refs,
        "source_file": path.name,
        "fingerprint": hashlib.sha256(json.dumps(cleaned[:200], sort_keys=True).encode()).hexdigest()[:16],
    }
    return cleaned, summary


def save_jsonl(records: list[dict], path: pathlib.Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def load_jsonl(path: pathlib.Path) -> list[dict]:
    return [json.loads(l) for l in path.read_text().splitlines() if l.strip()]