#!/usr/bin/env python3
import os
import sys
import torch
import pandas as pd

# Ensure routing_agent package is importable
ROOT = os.path.dirname(__file__)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from ppo_trainer import run_ppo_rollout, load_router_policy
from model import RouterNetwork
from sentence_transformers import SentenceTransformer


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Prepare embedder and model
    embed_model_name = 'all-MiniLM-L6-v2'
    embedder = SentenceTransformer(embed_model_name)

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

    # Optionally load existing checkpoint to continue training
    checkpoint_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint.pt')
    if os.path.exists(checkpoint_path):
        try:
            sd = torch.load(checkpoint_path, map_location='cpu')
            if isinstance(sd, dict) and 'model_state_dict' in sd:
                model.load_state_dict(sd['model_state_dict'], strict=False)
            else:
                model.load_state_dict(sd, strict=False)
            print('Loaded checkpoint', checkpoint_path)
        except Exception as e:
            print('Failed loading checkpoint:', e)

    optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)

    # training_queries = [
    #     "What are the graduation requirements for the 2026 academic year?",
    #     "Show me the syllabus for the Advanced Database course.",
    #     "How do I apply for the on-campus dormitory?"
    # ]
    df = pd.read_csv('df_queries.csv')
    training_queries = df['question'].tolist()
    
    # training_queries = training_queries[:2]  # Limit to first 200 queries for faster testing


    # Run PPO rollouts (this uses simulate_rag_pipeline to call the agent tools)
    run_ppo_rollout(model, embedder, training_queries, optimizer, epochs=40, device=str(device))

    # Save the updated model
    out_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint_pt_after_ppo.pt')
    torch.save({'model_state_dict': model.state_dict()}, out_path)
    print('Saved PPO-updated checkpoint to', out_path)


if __name__ == '__main__':
    main()