"""
Minimal Training Script using Unsloth & TRL for the OpenEnv Supply Chain Environment.
You can run this in a Google Colab notebook instance.
"""

# Install dependencies in Colab:
# !pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
# !pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes
# !pip install openenv-core

import torch
from unsloth import FastLanguageModel
from trl import SFTTrainer, PPOConfig, PPOTrainer
from transformers import TrainingArguments
from openenv_wrapper import SupplyChainEnv

# 1. Load Unsloth FastLanguageModel
max_seq_length = 2048
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = None,
    load_in_4bit = True,
)

# 2. Add LoRA adapters
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
)

# 3. Setup OpenEnv Environment
env = SupplyChainEnv(max_turns=12)

# For actual Reinforcement Learning (TRL + PPO), you would use the PPOTrainer 
# wrapped around the environment rollout loops.
# Here is a stub for the SFT/RL loop taking observations from the env:

def run_rl_episode(env, model, tokenizer):
    obs, info = env.reset()
    done = False
    
    # Simple prompt template for the model
    system_prompt = "You are an AI Supply Chain Manager. Respond ONLY with valid JSON actions."
    
    while not done:
        prompt = f"<|system|>\n{system_prompt}\n<|user|>\nObservation:\n{obs}\n<|assistant|>\n"
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        outputs = model.generate(**inputs, max_new_tokens=100)
        
        # Decode the action generated by the LLM
        action_str = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        
        # Step the environment
        obs, reward, done, truncated, info = env.step(action_str)
        print(f"Action Taken: {action_str} | Reward: {reward}")

# To train with SFT based on expert trajectories (if generated):
# trainer = SFTTrainer(
#     model = model,
#     train_dataset = your_expert_dataset,
#     dataset_text_field = "text",
#     max_seq_length = max_seq_length,
#     args = TrainingArguments(
#         per_device_train_batch_size = 2,
#         gradient_accumulation_steps = 4,
#         warmup_steps = 5,
#         max_steps = 60,
#         learning_rate = 2e-4,
#         fp16 = not torch.cuda.is_bfloat16_supported(),
#         bf16 = torch.cuda.is_bfloat16_supported(),
#         logging_steps = 1,
#         optim = "adamw_8bit",
#         output_dir = "outputs",
#     ),
# )
# trainer.train()

if __name__ == "__main__":
    print("Executing a test rollout with the untrained LoRA model...")
    run_rl_episode(env, model, tokenizer)