"""
ZDX Universal Training Data Formatter
=====================================
Convert mixed datasets into clean model-training corpora.

Project: https://zerodrivex.com/tools
Version: 1.0.1

Copyright (c) 2026 ZeroDriveX
SPDX-License-Identifier: MIT

Requirements:
    pip install chardet

Supported input formats:
    .txt        Plain text — one doc per file or one entry per line
    .json       JSON array or object — auto-detects text fields
    .jsonl      JSON Lines — one object per line (HuggingFace, OpenAI format)
    .csv        CSV — auto-detects text columns
    .tsv        TSV — same as CSV with tab delimiter
    .md         Markdown — strips formatting, keeps content

Output:
    A single clean corpus.txt ready for model training.
    One chunk per line, sized for the model's context window.

Usage:
    # Format a directory of mixed files
    python3 formatter.py --input ./raw_data --output ./corpus/corpus.txt

    # Format a single file
    python3 formatter.py --input ./data.jsonl --output ./corpus/corpus.txt

    # Specify JSON field manually
    python3 formatter.py --input ./data.json --output ./corpus/corpus.txt --field text

    # Set min/max chunk length
    python3 formatter.py --input ./raw --output ./corpus/corpus.txt --min_len 64 --max_len 512
"""

import os
import re
import csv
import json
import argparse
import chardet
from pathlib import Path


# ---------------------------------------------------------------------------
# Config defaults
# ---------------------------------------------------------------------------
MIN_CHUNK_LEN  = 64    # discard chunks shorter than this (characters)
MAX_CHUNK_LEN  = 1024  # split chunks longer than this
MIN_WORD_COUNT = 8     # discard chunks with fewer words than this

# Common JSON field names that usually contain the training text
# Checked in order — first match wins
JSON_TEXT_FIELDS = [
    "text", "content", "body", "article", "paragraph",
    "passage", "document", "sentence", "input", "output",
    "prompt", "completion", "message", "description", "abstract",
    "caption", "summary", "instruction", "response", "story",
]


# ---------------------------------------------------------------------------
# Encoding detection — handles files from anywhere
# ---------------------------------------------------------------------------
def read_file_bytes(path: str) -> bytes:
    with open(path, "rb") as f:
        return f.read()


def decode_bytes(raw: bytes) -> str:
    # Try UTF-8 first (fastest)
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError:
        pass
    # Detect encoding
    detected = chardet.detect(raw[:10000])
    enc = detected.get("encoding") or "latin-1"
    return raw.decode(enc, errors="replace")


# ---------------------------------------------------------------------------
# Text cleaning — normalises whatever comes in
# ---------------------------------------------------------------------------
def clean_text(text: str) -> str:
    if not text or not isinstance(text, str):
        return ""

    # Strip HTML tags
    text = re.sub(r"<[^>]+>", " ", text)

    # Strip markdown formatting but keep content
    text = re.sub(r"#{1,6}\s+", "", text)          # headers
    text = re.sub(r"\*{1,3}([^*]+)\*{1,3}", r"\1", text)  # bold/italic
    text = re.sub(r"`{1,3}[^`]*`{1,3}", " ", text) # code blocks
    text = re.sub(r"!\[.*?\]\(.*?\)", " ", text)   # images
    text = re.sub(r"\[([^\]]+)\]\([^\)]+\)", r"\1", text)  # links -> text

    # Normalise whitespace
    text = re.sub(r"\r\n|\r", "\n", text)
    text = re.sub(r"\t", " ", text)
    text = re.sub(r" {2,}", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)

    # Strip leading/trailing whitespace
    text = text.strip()

    return text


# ---------------------------------------------------------------------------
# Chunking — splits long text into model-sized pieces
# ---------------------------------------------------------------------------
def chunk_text(text: str, max_len: int = MAX_CHUNK_LEN, min_len: int = MIN_CHUNK_LEN) -> list[str]:
    chunks = []

    # First split on double newlines (paragraph boundaries)
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]

    for para in paragraphs:
        if len(para) <= max_len:
            if len(para) >= min_len:
                chunks.append(para)
        else:
            # Split long paragraph on sentence boundaries
            sentences = re.split(r"(?<=[.!?])\s+", para)
            current = ""
            for sent in sentences:
                if len(current) + len(sent) + 1 <= max_len:
                    current = (current + " " + sent).strip()
                else:
                    if len(current) >= min_len:
                        chunks.append(current)
                    current = sent
            if len(current) >= min_len:
                chunks.append(current)

    return chunks


def is_valid_chunk(text: str, min_words: int = MIN_WORD_COUNT) -> bool:
    """Filter out garbage chunks — too short, too many special chars, not real text."""
    words = text.split()
    if len(words) < min_words:
        return False
    # Reject if >30% non-alphanumeric (likely code garbage or binary noise)
    alpha = sum(1 for c in text if c.isalnum() or c.isspace())
    if alpha / len(text) < 0.70:
        return False
    return True


# ---------------------------------------------------------------------------
# Format parsers
# ---------------------------------------------------------------------------
def detect_json_field(obj: dict, user_field: str = None) -> str | None:
    """Find which field in a JSON object contains the training text."""
    if user_field and user_field in obj:
        return user_field

    # Check known field names
    for field in JSON_TEXT_FIELDS:
        if field in obj and isinstance(obj[field], str) and len(obj[field]) > 20:
            return field

    # Fall back: find the longest string field
    best_field = None
    best_len = 0
    for k, v in obj.items():
        if isinstance(v, str) and len(v) > best_len:
            best_len = len(v)
            best_field = k

    return best_field


def parse_txt(raw: str) -> list[str]:
    # Try line-by-line first
    lines = [l.strip() for l in raw.splitlines() if l.strip()]
    # If average line is short, treat whole file as one document
    avg_len = sum(len(l) for l in lines) / max(len(lines), 1)
    if avg_len < 100:
        return [raw]  # one big doc
    return lines


def parse_json(raw: str, user_field: str = None) -> list[str]:
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        print(f"  [Warning] JSON parse error: {e}")
        return []

    texts = []

    if isinstance(data, list):
        for item in data:
            if isinstance(item, str):
                texts.append(item)
            elif isinstance(item, dict):
                field = detect_json_field(item, user_field)
                if field:
                    texts.append(item[field])

    elif isinstance(data, dict):
        # Could be {"data": [...]} or {"train": [...]} wrapper
        for v in data.values():
            if isinstance(v, list):
                for item in v:
                    if isinstance(item, str):
                        texts.append(item)
                    elif isinstance(item, dict):
                        field = detect_json_field(item, user_field)
                        if field:
                            texts.append(item[field])
                break  # take first list found
            elif isinstance(v, str) and len(v) > 50:
                texts.append(v)

    return texts


def parse_jsonl(raw: str, user_field: str = None) -> list[str]:
    texts = []
    detected_field = user_field

    for i, line in enumerate(raw.splitlines()):
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
            if isinstance(obj, str):
                texts.append(obj)
            elif isinstance(obj, dict):
                if detected_field is None:
                    detected_field = detect_json_field(obj, user_field)
                    if detected_field:
                        print(f"  [JSONL] Auto-detected text field: '{detected_field}'")
                if detected_field and detected_field in obj:
                    texts.append(str(obj[detected_field]))
        except json.JSONDecodeError:
            continue

    return texts


def parse_csv(raw: str, delimiter: str = ",", user_field: str = None) -> list[str]:
    texts = []
    reader = csv.DictReader(raw.splitlines(), delimiter=delimiter)
    detected_field = user_field

    for row in reader:
        if detected_field is None:
            # Find best text column
            detected_field = detect_json_field(dict(row), user_field)
            if detected_field:
                print(f"  [CSV] Auto-detected text column: '{detected_field}'")

        if detected_field and detected_field in row:
            texts.append(row[detected_field])
        else:
            # Fall back: concatenate all string values
            texts.append(" ".join(str(v) for v in row.values() if v))

    return texts


def parse_markdown(raw: str) -> list[str]:
    # Remove code blocks entirely (not useful for general LM training)
    raw = re.sub(r"```[\s\S]*?```", "", raw)
    raw = re.sub(r"`[^`]+`", "", raw)
    return [raw]


# ---------------------------------------------------------------------------
# File dispatcher
# ---------------------------------------------------------------------------
def parse_file(path: str, user_field: str = None) -> list[str]:
    ext = Path(path).suffix.lower()
    raw_bytes = read_file_bytes(path)
    raw = decode_bytes(raw_bytes)

    if ext in (".txt",):
        return parse_txt(raw)
    elif ext in (".json",):
        return parse_json(raw, user_field)
    elif ext in (".jsonl", ".ndjson"):
        return parse_jsonl(raw, user_field)
    elif ext in (".csv",):
        return parse_csv(raw, delimiter=",", user_field=user_field)
    elif ext in (".tsv",):
        return parse_csv(raw, delimiter="\t", user_field=user_field)
    elif ext in (".md", ".markdown"):
        return parse_markdown(raw)
    else:
        # Try plain text as fallback
        print(f"  [Warning] Unknown extension '{ext}', treating as plain text")
        return parse_txt(raw)


# ---------------------------------------------------------------------------
# Main formatter
# ---------------------------------------------------------------------------
def format_corpus(
    input_path: str,
    output_path: str,
    user_field: str = None,
    min_len: int = MIN_CHUNK_LEN,
    max_len: int = MAX_CHUNK_LEN,
    min_words: int = MIN_WORD_COUNT,
    deduplicate: bool = True,
) -> dict:

    input_path = Path(input_path)
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Collect all files
    if input_path.is_dir():
        files = sorted([
            p for p in input_path.rglob("*")
            if p.suffix.lower() in (".txt", ".json", ".jsonl", ".ndjson", ".csv", ".tsv", ".md", ".markdown")
            and p.is_file()
        ])
    else:
        files = [input_path]

    print(f"[Formatter] Found {len(files)} file(s) to process")

    all_chunks = []
    seen = set()
    stats = {
        "files_processed": 0,
        "files_failed": 0,
        "raw_texts": 0,
        "chunks_kept": 0,
        "chunks_discarded": 0,
        "duplicates_removed": 0,
        "total_chars": 0,
    }

    for i, fpath in enumerate(files):
        print(f"  [{i+1}/{len(files)}] {fpath.name} ...", end=" ")
        try:
            raw_texts = parse_file(str(fpath), user_field)
            stats["raw_texts"] += len(raw_texts)

            file_chunks = 0
            for text in raw_texts:
                text = clean_text(text)
                if not text:
                    continue
                chunks = chunk_text(text, max_len=max_len, min_len=min_len)
                for chunk in chunks:
                    if not is_valid_chunk(chunk, min_words):
                        stats["chunks_discarded"] += 1
                        continue
                    if deduplicate:
                        key = chunk[:128]  # fingerprint first 128 chars
                        if key in seen:
                            stats["duplicates_removed"] += 1
                            continue
                        seen.add(key)
                    all_chunks.append(chunk)
                    stats["total_chars"] += len(chunk)
                    file_chunks += 1

            print(f"{file_chunks} chunks")
            stats["files_processed"] += 1

        except Exception as e:
            print(f"FAILED — {e}")
            stats["files_failed"] += 1

    stats["chunks_kept"] = len(all_chunks)

    # Write output
    with open(output_path, "w", encoding="utf-8") as f:
        for chunk in all_chunks:
            f.write(chunk.replace("\n", " ") + "\n")

    # Print report
    print(f"\n{'='*50}")
    print(f"[Formatter] Complete")
    print(f"  Files processed : {stats['files_processed']}")
    print(f"  Files failed    : {stats['files_failed']}")
    print(f"  Raw texts in    : {stats['raw_texts']:,}")
    print(f"  Chunks kept     : {stats['chunks_kept']:,}")
    print(f"  Chunks discarded: {stats['chunks_discarded']:,}")
    print(f"  Duplicates rm'd : {stats['duplicates_removed']:,}")
    print(f"  Total chars     : {stats['total_chars']:,}")
    print(f"  Avg chunk len   : {stats['total_chars'] // max(stats['chunks_kept'], 1)}")
    print(f"  Output          : {output_path}")
    print(f"{'='*50}")

    return stats


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="ZDX Training Data Formatter")
    parser.add_argument("--input",    required=True, help="Input file or directory")
    parser.add_argument("--output",   required=True, help="Output .txt file path")
    parser.add_argument("--field",    default=None,  help="JSON field name containing text (auto-detected if not set)")
    parser.add_argument("--min_len",  type=int, default=MIN_CHUNK_LEN,  help="Min chunk length in characters")
    parser.add_argument("--max_len",  type=int, default=MAX_CHUNK_LEN,  help="Max chunk length in characters")
    parser.add_argument("--min_words",type=int, default=MIN_WORD_COUNT, help="Min word count per chunk")
    parser.add_argument("--no_dedup", action="store_true", help="Disable deduplication")
    args = parser.parse_args()

    format_corpus(
        input_path=args.input,
        output_path=args.output,
        user_field=args.field,
        min_len=args.min_len,
        max_len=args.max_len,
        min_words=args.min_words,
        deduplicate=not args.no_dedup,
    )


if __name__ == "__main__":
    main()