import json
import random
from server.openenv_wrapper import SupplyChainEnv
from server.engine import CONTAINER_SIZE, SHIPPING_COST

def generate_expert_trajectories(num_episodes=20, output_file="dataset.jsonl"):
    env = SupplyChainEnv(max_turns=24)
    dataset = []

    print(f"Generating {num_episodes} expert episodes...")
    
    for episode in range(num_episodes):
        obs_json, info = env.reset()
        done = False
        
        while not done:
            obs = json.loads(obs_json)
            state_dict = env._state # Accessing internal state for heuristic planning
            
            # Find available ports
            open_ports = [p for p in state_dict.portDelays if not p["isClosed"]]
            
            if not open_ports:
                # No open ports, only possible action is to advance
                action_dict = {"action": "advance"}
                
                # Log the pair
                dataset.append({
                    "prompt": f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: ",
                    "completion": json.dumps(action_dict)
                })
                
                obs_json, reward, done, truncated, info = env.step(json.dumps(action_dict))
                continue
                
            # Best port and factory heuristic
            best_port = min(open_ports, key=lambda x: x["currentDelay"])
            best_factory = min(state_dict.factoryCosts, key=lambda x: x["currentCost"])
            
            actions_taken_this_turn = False
            
            for region, demand in state_dict.demands.items():
                if state_dict.inventory[region] < demand:
                    needed = demand - state_dict.inventory[region]
                    
                    cost_of_goods = best_factory["currentCost"] * needed
                    containers = (needed + CONTAINER_SIZE - 1) // CONTAINER_SIZE
                    total_cost = cost_of_goods + (containers * SHIPPING_COST)
                    
                    if state_dict.balance >= total_cost:
                        action_dict = {
                            "action": "order",
                            "factory": best_factory["id"],
                            "quantity": needed,
                            "port": best_port["id"],
                            "dest": region
                        }
                        
                        # Log the order action
                        dataset.append({
                            "prompt": f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: ",
                            "completion": json.dumps(action_dict)
                        })
                        
                        obs_json, reward, done, truncated, info = env.step(json.dumps(action_dict))
                        actions_taken_this_turn = True
            
            # After fulfilling all possible demands, we must advance the turn
            advance_action = {"action": "advance"}
            dataset.append({
                "prompt": f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: ",
                "completion": json.dumps(advance_action)
            })
            obs_json, reward, done, truncated, info = env.step(json.dumps(advance_action))
            
    # Write to JSONL
    with open(output_file, "w") as f:
        for entry in dataset:
            f.write(json.dumps(entry) + "\n")
            
    print(f"✅ Generated {len(dataset)} expert action pairs and saved to {output_file}")

if __name__ == "__main__":
    generate_expert_trajectories()