import sys
import os
import signal
import argparse
import uuid
import datetime

import typing
from typing import Union, List, Dict, Any, Optional

import gym
import gym_donkeycar

from stable_baselines import PPO2
# from stable_baselines.common import set_global_seeds
from stable_baselines.common.policies import CnnPolicy
from stable_baselines.common.vec_env import DummyVecEnv, VecEnv
from stable_baselines.common.callbacks import EventCallback
from stable_baselines.common.base_class import BaseRLModel

class MyCallback(EventCallback):
    def __init__(self, 
                 eval_env: Union[gym.Env, VecEnv],
                 save_freq: int = 1000,
                 save_file: Optional[str] = None,
                 verbose: int = 1):
        super(MyCallback, self).__init__(verbose=verbose)
        self.save_file = save_file
        self.save_freq = save_freq
        
        # Convert to VecEnv for consistency
        if not isinstance(eval_env, VecEnv):
            eval_env = DummyVecEnv([lambda: eval_env])
            
        assert eval_env.num_envs == 1, "You must pass only one environment for evaluation"
        
        self.eval_env = eval_env
    
    def init_callback(self, model: 'BaseRLModel') -> None:
        print("#### INIT ####")
        super(EventCallback, self).init_callback(model)
    
    def _init_callback(self):
        print("#### _INIT ####")
    
    def _on_step(self) -> bool:
        if self.save_freq > 0 and self.n_calls % self.save_freq == 0:
            now = datetime.datetime.now()
            if self.verbose > 0:
                print(f'{now} saving...')
            if self.save_file :
                now_str  = now.strftime('%y%m%d_%H%M%S')
                filename = self.save_file
                # filename = os.path.join(os.path.dirname(self.save_file), f'{now_str}_{os.path.basename(self.save_file)}')
                self.model.save(filename)
        return True
    

if __name__ == "__main__":
    # Initialize the donkey environment
    # where env_name one of:
    env_list = [
        "donkey-warehouse-v0",
        "donkey-generated-roads-v0",
        "donkey-avc-sparkfun-v0",
        "donkey-generated-track-v0",
        "donkey-roboracingleague-track-v0",
        "donkey-waveshare-v0",
        "donkey-minimonaco-track-v0",
        "donkey-warren-track-v0",
        "donkey-thunderhill-track-v0",
        "donkey-circuit-launch-track-v0",
    ]
    
    parser = argparse.ArgumentParser(description="ppo_train")
    parser.add_argument(
        "--sim",
        type=str,
        default="sim_path",
        help="path to unity simulator. maybe be left at manual if you would like to start the sim on your own.",
    )
    parser.add_argument("--model",     type=str, default="ppo_donkey",  help="path to model")
    parser.add_argument("--host",      type=str, default="127.0.0.1",   help="simulator address")
    parser.add_argument("--port",      type=int, default=9091,          help="port to use for tcp")
    parser.add_argument("--step",      type=int, default=10000,         help="port to use for tcp")
    parser.add_argument("--test_step", type=int, default=0,             help="port to use for tcp")
    parser.add_argument("--test",      action="store_true",             help="load the trained model and play")
    parser.add_argument("--env_name",  type=str, default="donkey-warehouse-v0",     help="name of donkey sim environment", choices=env_list)
    
    args = parser.parse_args()
    
    test_step = args.test_step
    if args.test and test_step == 0 :
        test_step = 1000
        
    # Complement the file extension
    if not args.model.endswith(".zip") :
        model_path = args.model + ".zip"
    else :
        model_path = args.model
    
    conf = {
        "exe_path": args.sim,
        "host": args.host,
        "port": args.port,
        "body_style": "car01",
        "body_rgb": (128, 255, 128),
        "car_name": "me",
        "font_size": 100,
        "racer_name": "PPO",
        "country": "USA",
        "bio": "Learning to drive w PPO RL",
        "guid": str(uuid.uuid4()),
        "max_cte": 10,
    }
    
    # Make an environment test our trained policy
    env = gym.make(args.env_name, conf=conf)
    env = DummyVecEnv([lambda: env])
    
    # hook terninate signal
    def signal_handler(signal, frame):
        print("catching ctrl+c")
        env.close()
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGABRT, signal_handler)
    
    try :
        # check model path
        if os.path.exists(model_path):
            # load model
            model = PPO2.load(model_path, env = env)
        else : 
            if not args.test:
                # create model
                print("create new model")
                model = PPO2(CnnPolicy, env, verbose=1)
            else :
                raise ValueError(f"Error: the file {model_path} could not be found")
        
        # change throttle lower limit
        env.action_space.low[1] = 0.1
        
        if not args.test:
            # in training mode
            
            callback = MyCallback(env, save_freq = 5000, save_file = model_path, verbose = 1)
            # callback = MyCallback(env, save_freq = 10, verbose = 1)
            
            # set up model in learning mode with goal number of timesteps to complete
            # model.learn(total_timesteps=10000)
            model.learn(total_timesteps=args.step, callback=callback)
            
            # save model
            model.save(model_path)
        
        print("stert testing...")
        obs = env.reset()
        prev_done_count = 0
        for i in range(test_step):
            action, _states = model.predict(obs)
            obs, rewards, dones, info = env.step(action)
            # print(f"cnt : {i}    rewards : {rewards[0]}    dones : {dones[0]}    pos : {info[0]['pos']}, CrossTrackError : {info[0]['cte']}, speed : {info[0]['speed']}")
            # print(f"+++ info: {info} +++")
            env.render()
            if (dones) :
                print(f'dones flag detected : {i - prev_done_count}  ({i})')
                prev_done_count = i
        print("done testing")
        
    except KeyboardInterrupt:
        print("stopping run...")
    
    env.close()