#!/usr/bin/env python3
"""
scripts/clean_prepared_data.py

Cleans the outputs in `prepared_data/` using heuristic filters to remove noisy / unwanted examples.

Produces:
- prepared_data/cleaned_extracted_text.jsonl
- prepared_data/cleaned_labeled_dataset.csv
- prepared_data/cleaned_instruction_tuning.jsonl
- prepared_data/cleaning_report.json

Usage examples (run from project root):

# Dry-run (no writes), print summary
python scripts/clean_prepared_data.py --prepared-dir prepared_data --dry-run

# Run and write cleaned outputs with defaults
python scripts/clean_prepared_data.py --prepared-dir prepared_data --apply

# More aggressive: drop pages with >10 BTC addresses or >30 [URL] markers
python scripts/clean_prepared_data.py --prepared-dir prepared_data --apply --url-threshold 30 --btc-threshold 10

The heuristics are intentionally conservative; tweak thresholds as needed.
"""

from pathlib import Path
import argparse
import json
import re
import hashlib
from collections import Counter, defaultdict
import csv
import sys

# Regexes
BTC_RE = re.compile(r"\b[13][a-km-zA-HJ-NP-Z1-9]{25,34}\b")
HEX_RE = re.compile(r"\b0x[a-fA-F0-9]{10,}\b")
LONG_HEX_RE = re.compile(r"\b[A-Fa-f0-9]{40,}\b")
URL_TOKEN = "[URL]"


def sha256_text(s: str) -> str:
    return hashlib.sha256(s.encode('utf-8', errors='ignore')).hexdigest()


def load_jsonl(path: Path):
    with path.open(encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except Exception:
                # try to be tolerant to bad lines
                continue


def write_jsonl(path: Path, objs):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open('w', encoding='utf-8') as f:
        for o in objs:
            f.write(json.dumps(o, ensure_ascii=False) + "\n")


def read_csv_as_rows(path: Path):
    with path.open(encoding='utf-8', newline='') as f:
        reader = csv.DictReader(f)
        for r in reader:
            yield r


def write_csv(path: Path, rows, fieldnames):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open('w', encoding='utf-8', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for r in rows:
            writer.writerow(r)


class Cleaner:
    def __init__(self, url_threshold=20, min_chars=60, btc_threshold=8, hex_threshold=3, repeat_chunk_thresh=3, blacklist=None, blacklist_enabled=False, max_chars=20000):
        self.url_threshold = int(url_threshold)
        self.min_chars = int(min_chars)
        self.btc_threshold = int(btc_threshold)
        self.hex_threshold = int(hex_threshold)
        self.repeat_chunk_thresh = int(repeat_chunk_thresh)
        self.blacklist = set([b.lower() for b in (blacklist or [])])
        self.blacklist_enabled = blacklist_enabled
        self.max_chars = int(max_chars)

    def _count_repeated_chunks(self, text, chunk_size=60):
        # simple repeated-chunk heuristic: count duplicate substrings of length chunk_size
        if len(text) < chunk_size * 2:
            return 0
        chunks = [text[i:i+chunk_size] for i in range(0, len(text)-chunk_size+1, chunk_size)]
        c = Counter(chunks)
        # how many chunks repeat more than once
        repeats = sum(1 for v in c.values() if v > 1)
        return repeats

    def is_noisy(self, text: str) -> (bool, str):
        """Return (is_noisy, reason) using heuristics observed in your dataset."""
        if not text or not text.strip():
            return True, 'empty'
        t = text
        # very short
        if len(t) < self.min_chars:
            return True, f'too_short:{len(t)}'
        # too many [URL] tokens (common in References blocks)
        url_tokens = t.count(URL_TOKEN) + len(re.findall(r'https?://', t)) + t.count('file://')
        if url_tokens >= self.url_threshold:
            return True, f'tool_many_urls:{url_tokens}'
        # btc address dumps
        btc_count = len(BTC_RE.findall(t))
        if btc_count >= self.btc_threshold:
            return True, f'btc_dump:{btc_count}'
        # hex-like tokens
        hex_count = len(HEX_RE.findall(t)) + len(LONG_HEX_RE.findall(t))
        if hex_count >= self.hex_threshold:
            return True, f'hex_tokens:{hex_count}'
        # long reference lists by word heuristics
        if 'compteur de visite' in t.lower() or ('references' in t.lower() and url_tokens > 5):
            return True, 'references_block'
        # repeated chunk heuristic
        repeats = self._count_repeated_chunks(t, chunk_size=80)
        if repeats >= self.repeat_chunk_thresh:
            return True, f'repeated_chunks:{repeats}'
        # blacklist
        if self.blacklist_enabled:
            low = t.lower()
            for b in self.blacklist:
                if b and b in low:
                    return True, f'blacklist_match:{b}'
        # too long single example: truncate instead of drop
        if len(t) > self.max_chars:
            # not noisy per se, but will be truncated by the cleaner
            return False, 'too_long_truncate'
        return False, ''

    def clean_text(self, text: str) -> str:
        # basic post-processing: collapse whitespace and truncate to max_chars
        s = re.sub(r"\s+", " ", text).strip()
        if len(s) > self.max_chars:
            s = s[:self.max_chars] + "\n\n[TRUNCATED]"
        return s


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--prepared-dir', type=Path, default=Path('prepared_data'))
    p.add_argument('--apply', action='store_true', help='Write cleaned outputs. Without this flag runs a dry-run and prints summary')
    p.add_argument('--url-threshold', type=int, default=20)
    p.add_argument('--min-chars', type=int, default=60)
    p.add_argument('--btc-threshold', type=int, default=8)
    p.add_argument('--hex-threshold', type=int, default=3)
    p.add_argument('--repeat-chunk-thresh', type=int, default=3)
    p.add_argument('--blacklist', type=str, default='', help='Comma-separated blacklist terms to drop (optional)')
    p.add_argument('--enable-blacklist', action='store_true', help='Enable blacklist matching')
    p.add_argument('--max-chars', type=int, default=20000)
    args = p.parse_args()

    prepared = Path(args.prepared_dir)
    if not prepared.exists():
        print('prepared_data dir not found:', prepared)
        sys.exit(2)

    extracted_path = prepared / 'extracted_text.jsonl'
    labeled_csv = prepared / 'labeled_dataset.csv'
    instruction_j = prepared / 'instruction_tuning.jsonl'

    if not extracted_path.exists():
        print('extracted_text.jsonl missing at', extracted_path)
        sys.exit(2)

    cleaner = Cleaner(
        url_threshold=args.url_threshold,
        min_chars=args.min_chars,
        btc_threshold=args.btc_threshold,
        hex_threshold=args.hex_threshold,
        repeat_chunk_thresh=args.repeat_chunk_thresh,
        blacklist=[t.strip() for t in args.blacklist.split(',')] if args.blacklist else None,
        blacklist_enabled=args.enable_blacklist,
        max_chars=args.max_chars,
    )

    cleaned_objs = []
    stats = defaultdict(int)
    reasons = Counter()
    kept_sources = set()
    seen_hashes = set()

    # iterate and filter
    for obj in load_jsonl(extracted_path):
        text = obj.get('text','')
        is_noise, reason = cleaner.is_noisy(text)
        if is_noise:
            stats['dropped'] += 1
            reasons[reason] += 1
            continue
        # clean text
        new_text = cleaner.clean_text(text)
        # dedupe by text hash
        h = sha256_text(new_text)[:32]
        if h in seen_hashes:
            stats['duplicate'] += 1
            reasons['duplicate'] += 1
            continue
        seen_hashes.add(h)
        # update object
        obj['text'] = new_text
        cleaned_objs.append(obj)
        kept_sources.add(obj.get('source_path') or obj.get('filename') or '')
        stats['kept'] += 1

    stats['total'] = stats.get('kept',0) + stats.get('dropped',0) + stats.get('duplicate',0)

    # prepare output paths
    out_extracted = prepared / 'cleaned_extracted_text.jsonl'
    out_csv = prepared / 'cleaned_labeled_dataset.csv'
    out_instr = prepared / 'cleaned_instruction_tuning.jsonl'
    report_path = prepared / 'cleaning_report.json'

    # dry-run: print summary
    if not args.apply:
        print('DRY RUN - no files written')
        print('Prepared dir:', prepared)
        print('Total input examples:', stats['total'])
        print('Kept:', stats['kept'], 'Dropped:', stats['dropped'], 'Duplicates:', stats['duplicate'])
        print('Top drop reasons:')
        for r, c in reasons.most_common(12):
            print(' ', r, c)
        print('\nTo write cleaned outputs add --apply')
        return

    # write cleaned extracted JSONL
    write_jsonl(out_extracted, cleaned_objs)

    # filter labeled CSV to keep only rows that are in cleaned set
    if labeled_csv.exists():
        kept_rows = []
        fieldnames = None
        for row in read_csv_as_rows(labeled_csv):
            # match by source_path
            sp = row.get('source_path') or ''
            if sp in kept_sources:
                kept_rows.append(row)
            else:
                # there may be entries where source_path is just a filename; try filename match
                fn = row.get('source_path','')
        if kept_rows:
            fieldnames = list(kept_rows[0].keys())
            write_csv(out_csv, kept_rows, fieldnames)
        else:
            # fallback: regenerate labeled CSV minimally from cleaned_objs using Unknown label
            fieldnames = ['source_path','text','label']
            regen = []
            for o in cleaned_objs:
                regen.append({'source_path': o.get('source_path',''), 'text': o.get('text',''), 'label': 'Unknown'})
            write_csv(out_csv, regen, fieldnames)
    else:
        # No labeled CSV original: create minimal one
        fieldnames = ['source_path','text','label']
        regen = []
        for o in cleaned_objs:
            regen.append({'source_path': o.get('source_path',''), 'text': o.get('text',''), 'label': 'Unknown'})
        write_csv(out_csv, regen, fieldnames)

    # regenerate instruction JSONL from cleaned CSV
    # Use same basic prompt as original pipeline to keep format consistent
    PROMPT = (
        """Below is an instruction that describes a task, paired with an input that provides further context. \nWrite a response that appropriately completes the request. \nBefore answering, think carefully about the question ensure an accurate response.\n\n### Instruction:\nYou are an expert in cybersecurity and threat intelligence.\nYour role is to provide precise classification of threats from the description provided by the user.\n\n### Description : \n{} \n\n### Response:\n{}"""
    )

    written = 0
    with out_instr.open('w', encoding='utf-8') as outf:
        # read from cleaned CSV
        for row in read_csv_as_rows(out_csv):
            text = row.get('text','')
            label = row.get('label','Unknown')
            prompt = PROMPT.format(text, '')
            combined = prompt + "\n" + label + " <EOS>"
            outf.write(json.dumps({'text': combined, 'label': label, 'source': row.get('source_path','')}, ensure_ascii=False) + "\n")
            written += 1

    # write report
    report = {
        'input_total': stats['total'],
        'kept': stats['kept'],
        'dropped': stats['dropped'],
        'duplicates': stats['duplicate'],
        'drop_reasons': dict(reasons.most_common()),
        'out_extracted': str(out_extracted),
        'out_csv': str(out_csv),
        'out_instruction': str(out_instr),
    }
    with report_path.open('w', encoding='utf-8') as f:
        json.dump(report, f, indent=2)

    print('Wrote cleaned extracted JSONL ->', out_extracted)
    print('Wrote cleaned labeled CSV ->', out_csv)
    print('Wrote cleaned instruction JSONL ->', out_instr)
    print('Wrote cleaning report ->', report_path)
    print('Summary: kept', stats['kept'], 'dropped', stats['dropped'], 'duplicates', stats['duplicate'])


if __name__ == '__main__':
    main()