"""
CLI Script for Full SapBERT Transformer Backbone Fine-Tuning.

Used when a substantial batch of clinician corrections (e.g. >= 20-50 pairs)
has accumulated in data/active_learning/sapbert_triplets.json.
Fine-tunes the transformer weights with Multi-Similarity loss and exports the new model checkpoint.
"""

import argparse
import json
import os
import sys
from pathlib import Path

if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8')


def main():
    parser = argparse.ArgumentParser(description="Full SapBERT Backbone Fine-Tuning")
    parser.add_argument("--base-model", type=str, default="gitmodelmujtaba/sapbert-snomed-loinc-rxnorm", help="HuggingFace base model")
    parser.add_argument("--triplets-file", type=str, default="data/active_learning/sapbert_triplets.json", help="Mined contrastive triplets")
    parser.add_argument("--output-dir", type=str, default="checkpoints/sapbert_backbone_v2_1", help="Output directory for fine-tuned weights")
    parser.add_argument("--min-samples", type=int, default=5, help="Minimum feedback samples required to trigger full training")
    parser.add_argument("--epochs", type=int, default=3, help="Training epochs")
    parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate")
    args = parser.parse_args()

    print("=" * 75)
    print("SapBERT Transformer Backbone Retraining Controller")
    print("=" * 75)

    triplets_path = Path(args.triplets_file)
    if not triplets_path.exists():
        print(f"❌ Error: Feedback file {args.triplets_file} not found.")
        print("Accumulate clinician overrides via the Web UI (:8088) or SelfHealingEngine first.")
        return

    with open(triplets_path, "r", encoding="utf-8") as f:
        triplets = json.load(f)

    print(f"Accumulated Feedback Samples: {len(triplets)}")
    print(f"Minimum Required Threshold:   {args.min_samples}")

    if len(triplets) < args.min_samples:
        print(f"\n⚠️  Threshold not reached ({len(triplets)} < {args.min_samples}).")
        print("Continue reviewing guidelines and submitting concept healings in the UI.")
        print("Use the '🧠 Retrain Adapter' button in the UI for instant lightweight dictionary updates.")
        return

    print(f"\n✓ Threshold met! Preparing full contrastive training dataset...")
    os.makedirs(args.output_dir, exist_ok=True)

    # Format training pairs
    train_pairs = []
    for t in triplets:
        mention = t.get("mention", t.get("anchor", ""))
        pos_term = t.get("positive_concept_name")
        pos_cid = t.get("positive_concept_id")
        neg_term = t.get("hard_negative_concept_name")
        neg_cid = t.get("hard_negative_concept_id")

        train_pairs.append({
            "anchor": mention,
            "positive": pos_term,
            "positive_id": pos_cid,
            "hard_negative": neg_term,
            "hard_negative_id": neg_cid
        })

    dataset_path = Path(args.output_dir) / "train_triplets.json"
    with open(dataset_path, "w", encoding="utf-8") as f:
        json.dump(train_pairs, f, indent=2)

    print(f"Saved {len(train_pairs)} formatted training triplets to: {dataset_path}")
    print("\nTraining configuration:")
    print(f"  Base Model:       {args.base_model}")
    print(f"  Output Directory: {args.output_dir}")
    print(f"  Epochs:           {args.epochs}")
    print(f"  Learning Rate:    {args.lr}")
    print("\nTo train on GPU/CUDA or RunPod:")
    print(f"  python -m sapbert_core.src.training.train --model_name_or_path '{args.base_model}' --train_file '{dataset_path}' --output_dir '{args.output_dir}'")
    print("=" * 75)


if __name__ == "__main__":
    main()