#!/usr/bin/env python3
import os
import sys
import argparse
import csv
import torch

ROOT = os.path.dirname(__file__)

def build_dataset_from_csv(csv_path):
    dataset = []
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            q = row.get('question') or row.get('Question')
            lbl = row.get('label') or row.get('Label')
            if q is None or lbl is None:
                continue
            try:
                lbl = int(lbl)
            except Exception:
                continue
            dataset.append({'question': q, 'label': lbl})
    return dataset


class EmbedderWrapper:
    """Wrap a SentenceTransformer instance to ensure tensors are on `device`.

    The wrapper exposes `encode(texts, convert_to_tensor=True, **kwargs)` same as
    the SentenceTransformer API and will move returned torch.Tensor to the device.
    """
    def __init__(self, embedder, device):
        self.embedder = embedder
        self.device = device

    def encode(self, texts, convert_to_tensor=True, **kwargs):
        emb = self.embedder.encode(texts, convert_to_tensor=convert_to_tensor, **kwargs)
        # Move tensor outputs to device (if it's a torch.Tensor)
        if convert_to_tensor and hasattr(emb, 'to'):
            try:
                emb = emb.to(self.device)
            except Exception:
                pass
        return emb


def main():
    parser = argparse.ArgumentParser(description='Run SFT training for RouterNetwork')
    parser.add_argument('--csv', default='processed_dataset.csv', help='Path to CSV dataset')
    parser.add_argument('--model_name', default='all-MiniLM-L6-v2', help='SentenceTransformer model name')
    parser.add_argument('--epochs', type=int, default=5)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--checkpoint', default='router_checkpoint.pt', help='Checkpoint output path')
    parser.add_argument('--device', default=None, help='Device to use (cpu or cuda). Auto-detect if unset')
    args = parser.parse_args()

    # Ensure routing_agent directory is importable such that its modules' imports work
    routing_agent_path = os.path.join(ROOT, 'routing_agent')
    if routing_agent_path not in sys.path:
        sys.path.insert(0, routing_agent_path)

    # Import modules from routing_agent (sft_trainer expects to import model as a top-level module)
    try:
        import sft_trainer
        import model as model_mod
    except Exception as e:
        print('Failed importing routing_agent modules:', e)
        sys.exit(1)

    # Device selection
    if args.device:
        device = torch.device(args.device)
    else:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Build dataset
    csv_path = os.path.join(ROOT, args.csv) if not os.path.isabs(args.csv) else args.csv
    if not os.path.exists(csv_path):
        print('Dataset CSV not found at', csv_path)
        sys.exit(1)

    dataset = build_dataset_from_csv(csv_path)
    if len(dataset) == 0:
        print('No valid rows found in dataset. Exiting.')
        sys.exit(1)

    # Initialize embedder
    try:
        from sentence_transformers import SentenceTransformer
    except Exception as e:
        print('sentence-transformers not installed or failed to import:', e)
        sys.exit(1)

    embedder = SentenceTransformer(args.model_name)
    embedder_wrapper = EmbedderWrapper(embedder, device)

    # Infer embedding dim and instantiate model from the same module sft_trainer expects
    try:
        embed_dim = embedder.get_sentence_embedding_dimension()
    except Exception:
        # fallback to common dim
        embed_dim = 384

    model = model_mod.RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)

    # Run SFT training
    print(f"Starting SFT training on device={device} with model={args.model_name} | epochs={args.epochs} | lr={args.lr}")
    sft_trainer.train_sft(model, embedder_wrapper, dataset, epochs=args.epochs, lr=args.lr)

    # Save checkpoint
    try:
        torch.save(model.state_dict(), args.checkpoint)
        print('Saved checkpoint to', args.checkpoint)
    except Exception as e:
        print('Failed to save checkpoint:', e)


if __name__ == '__main__':
    main()