import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


def main():
    parser = argparse.ArgumentParser(description="Classify a text log file one line at a time.")
    parser.add_argument("model", help="Local model directory or Hugging Face repo id")
    parser.add_argument("input", help="Input text log file")
    parser.add_argument("--output", default="classified.jsonl")
    parser.add_argument("--threshold", type=float, default=0.5)
    args = parser.parse_args()

    torch.set_num_threads(2)
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    model = AutoModelForSequenceClassification.from_pretrained(args.model, trust_remote_code=True)
    model.eval()

    input_path = Path(args.input)
    output_path = Path(args.output)

    with input_path.open("r", encoding="utf-8", errors="replace") as src, output_path.open("w", encoding="utf-8") as dst:
        for line_no, raw in enumerate(src, 1):
            text = raw.rstrip("\r\n")
            if not text:
                continue
            encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=96)
            with torch.inference_mode():
                probs = torch.softmax(model(**encoded).logits, dim=-1)[0]
            suspicious = float(probs[1])
            label = "SUSPICIOUS" if suspicious >= args.threshold else "BENIGN"
            dst.write(json.dumps({
                "line": line_no,
                "label": label,
                "suspicious_probability": round(suspicious, 6),
                "text": text,
            }) + "\n")

    print(f"Wrote {output_path}")


if __name__ == "__main__":
    main()