import pickle
from pathlib import Path
import warnings

from agentdriver.llm_core.timeout import timeout
from agentdriver.planning.planning_target import (
    generate_planning_target,
)
from agentdriver.planning.motion_planning import (
    planning_batch_inference,
    planning_single_inference,
)

class PlanningAgent:
    def __init__(self, model_name="", verbose=True) -> None:
        self.verbose = verbose
        self.model_name = model_name # Note: this model must be a **finetuned** GPT model
        if model_name == "" or model_name[:2] != "ft":
            warnings.warn(f"Input motion planning model might not be correct, \
                  expect a fintuned model like ft:gpt-3.5-turbo-0613:your_org::your_model_id, \
                  but get {self.model_name}", UserWarning)
    
    def generate_planning_target(self, data_dict):
        planning_target = generate_planning_target(data_dict)
        if self.verbose:
            print(planning_target)
        return planning_target

    @timeout(15)
    def generate_target(self, data_dict):
        """Generate planning target and chain_of_thoughts reasoning"""
        planning_target = self.generate_planning_target(data_dict)
        return planning_target
    
    @timeout(15)
    def run(self, data_dict, data_sample):
        """Generate motion planning results for a single scene"""
        planning_traj = planning_single_inference(
            planner_model_id=self.model_name, 
            data_sample = data_sample, 
            data_dict=data_dict, 
            verbose=self.verbose
        )
        save_file_name = "experiments/redteaming/pred_trajs_dict.pkl"
        token = data_sample["token"]
        pred_trajs_dict = {}
        pred_trajs_dict[token] = planning_traj[0]
        with open(save_file_name, "wb") as f:
            pickle.dump(pred_trajs_dict, f)
        return planning_traj

    def run_batch(self, data_samples, data_path, save_path, args=None):
        planning_traj_dict = planning_batch_inference(
            data_samples=data_samples, 
            planner_model_id=self.model_name, 
            data_path=data_path, 
            save_path=save_path,
            use_local_planner=False,
            args=args,
        )
        return planning_traj_dict