import torch
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer
from server.openenv_wrapper import SupplyChainEnv

def train_ppo():
    print("Loading SFT Model for PPO Training...")
    
    # Load the SFT model we just trained
    # Note: On a machine without a GPU, this will fail or be extremely slow.
    model = AutoModelForCausalLMWithValueHead.from_pretrained("tradestorm-sft-model")
    tokenizer = AutoTokenizer.from_pretrained("tradestorm-sft-model")

    ppo_config = PPOConfig(
        batch_size=16, 
        mini_batch_size=4, 
        learning_rate=1e-5
    )
    
    print("Initializing PPO Trainer...")
    ppo_trainer = PPOTrainer(config=ppo_config, model=model, tokenizer=tokenizer)

    print("Initializing Supply Chain OpenEnv Wrapper...")
    env = SupplyChainEnv(max_turns=24)

    # Training Loop (500 episodes for demonstration)
    print("Starting PPO Training Loop...")
    for epoch in range(500):
        obs_json, info = env.reset()
        done = False
        
        episode_reward = 0
        
        while not done:
            # 1. Format the Observation
            prompt = f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: "
            inputs = tokenizer(prompt, return_tensors="pt")
            
            # If using GPU, move tensors to CUDA
            if torch.cuda.is_available():
                inputs = {k: v.to("cuda") for k, v in inputs.items()}
            
            # 2. LLM Generates Action
            action_tensors = ppo_trainer.generate(**inputs, max_new_tokens=100)
            
            # Decode the action
            action_str = tokenizer.decode(action_tensors[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
            
            # 3. Step the Environment
            obs_json, reward, done, truncated, info = env.step(action_str)
            episode_reward += reward
            
            # 4. PPO Weight Update
            reward_tensor = torch.tensor([reward], dtype=torch.float32)
            if torch.cuda.is_available():
                reward_tensor = reward_tensor.to("cuda")
                
            ppo_trainer.step([inputs["input_ids"][0]], [action_tensors[0]], [reward_tensor])

        print(f"Epoch {epoch} finished. Total Reward: {episode_reward} | Final Balance: {info.get('balance', 0)}")

    print("Saving Reinforcement Learning Model...")
    model.save_pretrained("tradestorm-rl-agent")
    print("✅ PPO Training complete. Model saved to tradestorm-rl-agent/")

if __name__ == "__main__":
    train_ppo()