use std::process::{Command,Stdio};
use std::{thread, time, str};
use std::time::{Duration, Instant};
use std::cmp::{max, min};
use std::fs;
use std::fs::File;
use std::path::Path;
use std::string::String;
use std::io::prelude::*;
use std::io::stdout;
use std::net::{TcpStream, SocketAddr};
use std::sync::mpsc::channel;
use std::env;
use std::iter::repeat;
use std::sync::mpsc;
use std::sync::Arc;
use std::f64;
use rustc_serialize::json::{Json, Object};

extern crate glob;
use glob::glob;

extern crate rand;
extern crate x11;
extern crate vnc;
extern crate image;

extern crate rustc_serialize;
use rustc_serialize::json;

extern crate hyper;
use self::hyper::header::{Headers, Authorization, Basic};

extern crate websocket;
use self::websocket::{Message, Sender, Receiver, Client};
use self::websocket::client::request::Url;
use self::websocket::client;
use self::websocket::client::Request;
use self::websocket::header::extensions::Extension;

use std::collections::HashMap;

pub type gym_reward = f64;
pub type gym_done = bool;
pub type gym_range = Vec<usize>;
#[derive(Debug, Clone)]
pub struct GymShape {
   pub action_space: gym_range,
   pub observation_space: gym_range,
   pub reward_max: gym_reward,
   pub reward_min: gym_reward
}
#[derive(Debug, Clone)]
pub struct GymState {
   pub screen: Vec<u8>
}
pub trait GymMember {
   fn start (&mut self, s: &GymShape, t: &GymState) -> ();
   fn reward (&mut self, gym_reward, gym_done) -> ();
   fn reset (&mut self) -> ();
   fn tick (&mut self, &mut Vec<u8>) -> gym_range;
   fn close (&mut self) -> ();
}
pub struct GymRemote {
   black_screen: bool,
   dirty: bool,
   id: u32,
   fps: u32,
   env_id: String,
   duration: u64,
   record_dst: String,
   mode: String,
   vnc: Option<vnc::Client>,
   rewarder: Option<((self::websocket::client::Sender<websocket::WebSocketStream>, self::websocket::client::Receiver<websocket::WebSocketStream>))>,
   state: GymState,
   shape: GymShape,
   offset_width: usize,
   offset_height: usize,
   time: Instant,
   frame: u32
}

pub struct Gym {
   fps: u32,
   env_id: String,
   max_parallel: u32,
   duration: u64,
   record_dst: String
}

#[derive(RustcDecodable, RustcEncodable)]
struct RCBody {
   seed: u32,
   env_id: String,
   fps: u32
}

#[derive(RustcDecodable, RustcEncodable)]
struct RCHeaders {
   sent_at: u32,
   episode_id: u32,
   message_id: u32
}

#[derive(RustcDecodable, RustcEncodable)]
struct RewarderCommand {
   method: String,
   body: RCBody,
   headers: RCHeaders
}

const ATARI_HEIGHT: u32 = 262;
const ATARI_WIDTH: u32 = 160;



impl GymRemote {
   pub fn recorder_cleanup(&mut self) -> () {
      let real_fps = self.frame / max(self.time.elapsed().as_secs() as u32, 1);
      Command::new("ffmpeg")
                 .arg("-r").arg(format!("{}", real_fps))
                 .arg("-f").arg("image2")
                 .arg("-i").arg("mov_out/frame_%0d.png")
                 .arg("-r").arg("24")
                 .arg("-pix_fmt").arg("yuv420p")
                 .arg(self.record_dst.clone())
                 .spawn();
   }
   pub fn start<T: GymMember>(&mut self, mut agent: T) {
      let rx = self.start_rewarder();
      self.start_vnc();
      let mut agent = self.start_agent(agent);
      let mut prev_action = Vec::new();

      use x11::keysym::*;
      let atari_actions = &[XK_Left, XK_Right, XK_Up, XK_Down, XK_space];
      let flash_actions = &[XK_Pointer_Button1, XK_Left, XK_Right, XK_Up, XK_Down, XK_space];

      let mut high_score = 0.0;

      loop {
         if self.time.elapsed().as_secs() > self.duration { break; }

         // Sync Rewarder - unrolled here due to blocking issue in library
         loop {
            match rx.try_recv() {
               Ok(s) => {
                  let ms = &s[..];
                  let ms = Json::from_str(ms).unwrap();
                  let o = ms.as_object().unwrap();
                  let ms = o.get("method").unwrap();
                  let ms = ms.as_string().unwrap();
                  match ms {
                     "v0.env.describe" => { }
                     "v0.env.reward" => {
                        let headers = o.get("headers").unwrap();
                        let headers = headers.as_object().unwrap();
                        let sent_at = headers.get("sent_at").unwrap();
                        let sent_at = sent_at.as_f64().unwrap();
                        let message_id = headers.get("message_id").unwrap();
                        let message_id = message_id.as_u64().unwrap();
                        let body = o.get("body").unwrap();
                        let body = body.as_object().unwrap();
                        let reward = body.get("reward").unwrap();
                        let reward = reward.as_f64().unwrap();
                        let done = body.get("done").unwrap();
                        let done = done.as_boolean().unwrap();
                        high_score = high_score + reward;
                        agent.reward( reward, done );
                     }
                     "v0.reply.env.reset" => {
                        println!("finished game with score: {}", high_score);
                        high_score = 0.0;
                        agent.reset();
                     }
                     _ => { println!("received unknown rewarder message: {}", s); }
                  }
               }
               Err(_) => { break; }
            }
         } 
         self.sync();
         let mut screen_copy = self.state.screen.clone();
         let action = agent.tick(&mut screen_copy);
         self.render_frame(screen_copy);


         let mut vnc = self.vnc.as_mut().unwrap();
         if action == prev_action { //pass
         } else if self.mode=="atari" && action.len()==5 {
            for ai in 0..5 {
               if prev_action.len()<5 || prev_action[ai]==action[ai] {
               } else if action[ai]==0 {
                  vnc.send_key_event(false, atari_actions[ai]).unwrap();
               } else {
                  vnc.send_key_event(true, atari_actions[ai]).unwrap();
               }
            }
         } else if self.mode=="flash" && action.len()==8 {
            let x = action[0];
            let y = action[1];
            if prev_action.len()==0 || (x!=prev_action[0] && y!=prev_action[1]) {
               //TODO move cursor
            }
            for ai in 2..8 {
               if prev_action.len()<8 || prev_action[ai]==action[ai] {
               } else if action[ai]==0 {
                  vnc.send_key_event(false, flash_actions[ai-2]).unwrap();
               } else {
                  vnc.send_key_event(true, flash_actions[ai-2]).unwrap();
               }
            }
         }
         prev_action = action.clone();
      }

      println!("disrupted game with score: {}", high_score);
      self.recorder_cleanup();
   }
   pub fn sync(&mut self) {
      self.sync_rewarder();
      self.sync_vnc();
   }
   pub fn start_agent<T: GymMember>(&mut self, mut agent: T) -> T {
      agent.start(&self.shape, &self.state);
      self.time = Instant::now();
      return agent
   }
   pub fn start_rewarder(&mut self) -> std::sync::mpsc::Receiver<String> {
      let ws_url = &format!("ws://127.0.0.1:{}", 15900+self.id)[..];
      let url = Url::parse(ws_url).unwrap();
      println!("Connecting to rewarder at {}", url);
      let mut request = Client::connect(url).unwrap();

      let mut auth_header = self::hyper::header::Authorization(
         self::hyper::header::Basic {
            username: "openai".to_owned(),
            password: Some("openai".to_owned())
         }
      );

      request.headers.set(auth_header);
      let mut response = request.send().unwrap(); // Send the request and retrieve a response

      response.validate().unwrap(); // Validate the response
      let (mut sender, mut receiver) = response.begin().split();

      //ws library not updated to use nonblocking yet
      //sender.set_nonblocking(true); 
      //receiver.set_nonblocking(true);

      // websocket receiver nonblocking is bugged, so we'll just put it in it's own thread...
      println!("Recorder websocket is now ready to use at {}", 15900+ self.id); 

      let reset_cmd = RewarderCommand{
         method: "v0.env.reset".to_owned(),
         body: RCBody {
            seed: 1,
            env_id: self.env_id.trim().to_string(),
            fps: self.fps
         },
         headers: RCHeaders {
            sent_at: 0,
            episode_id: 0,
            message_id: 0
         }
      };
      let reset_msg = json::encode(&reset_cmd).unwrap();
      println!("Send reset json message to rewarder: {}", reset_msg);
      let rst_msg = Message::text(String::from(reset_msg));
      sender.send_message(&rst_msg);

      //self.rewarder = Some((sender,receiver));

      let (tx, rx) = channel();
      thread::spawn(move || {
         for msg in receiver.incoming_messages() {
            let msg: Result<websocket::message::Message,_> = msg;
            match msg {
               Ok(message) => {
                  match message.opcode {
                     websocket::message::Type::Close => {
                        //agent.close()
                        println!("Connection to rewarder terminated.");
                     },
                     websocket::message::Type::Ping => { 
                        let mut pong_message = Message::pong(message.payload);
                        sender.send_message(&pong_message);
                     },
                     websocket::message::Type::Text => {
                        let bytes = message.payload.into_owned();
                        let msg = String::from_utf8(bytes).unwrap();
                        tx.send(msg).unwrap();
                     },
                     _ => {}
                  }
               }
               Err(e) => {
                  //if no frames available, IOError occurs
               }
            }
         }
      });

      return rx;
   }
   pub fn start_vnc(&mut self) {
      //connect vnc
      let vnc_addr: SocketAddr = format!("127.0.0.1:{}", 5900+self.id).parse().expect("Unable to parse socket address");
      let stream = match std::net::TcpStream::connect(vnc_addr) {
         Ok(stream) => stream,
         Err(error) => {
            panic!("cannot connect to localhost:{}: {}", 5900+self.id, error);
            std::process::exit(1)
         }
      };
      self.vnc = Some(match vnc::Client::from_tcp_stream(stream, true, |methods| {
         for method in methods {
            match method {
               &vnc::client::AuthMethod::Password => {
                  let mut key = [0; 8];
                  for (i, byte) in "openai".bytes().enumerate() {
                     if i == 8 { break }
                     key[i] = byte
                  }
                  return Some(vnc::client::AuthChoice::Password(key))
               }
               _ => ()
            }
         }
         None
      }) {
         Ok(vnc) => vnc,
         Err(error) => {
            panic!("cannot initialize VNC session: {}", error);
            std::process::exit(1)
         }
      });
      println!("Connected to vnc on port: {}", 5900+self.id);
   }
   pub fn sync_rewarder(&mut self) -> () {
      //let mut sr = self.rewarder.as_mut().unwrap();
      //let &mut (ref mut sender, ref mut receiver) = sr;

      //requires nonblocking before moving into main thread
   }
   pub fn render_frame(&mut self, screen: Vec<u8>) {
      let width = self.shape.observation_space[0];
      let height = self.shape.observation_space[1];

      if self.dirty {
         let mut imgbuf = image::ImageBuffer::new( width as u32, height as u32 );

         for x in 0 .. width {
            for y in 0 .. height {
               let left = 3*(y * width + x) as usize;
               imgbuf.put_pixel(x as u32, y as u32, image::Rgb([ screen[left], screen[left+1], screen[left+2] ]));
            }
         }

         let ref mut fout = File::create(&Path::new( &format!("mov_out/frame_{}.png", self.frame)[..] )).unwrap();
         let _ = image::ImageRgb8(imgbuf).save(fout, image::PNG);
         self.dirty = false;
      } else {
         let from_p = format!("mov_out/frame_{}.png", self.frame-1);
         let to_p = format!("mov_out/frame_{}.png", self.frame);
         std::fs::copy(Path::new(&from_p), Path::new(&to_p));
      }

      self.frame = self.frame + 1;
   }
   pub fn sync_vnc(&mut self) -> () {
      let width = self.shape.observation_space[0];
      let height = self.shape.observation_space[1];

      let pause = 1000 / self.fps;
      std::thread::sleep_ms(pause);

      let mut vnc = self.vnc.as_mut().unwrap();

      vnc.request_update(vnc::Rect {
         left: self.offset_width as u16,
         top: self.offset_height as u16,
         width: (self.offset_width+width) as u16,
         height: (self.offset_height+height) as u16
      }, false).unwrap();
      for event in vnc.poll_iter() {
         use vnc::client::Event;
         match event {
            Event::EndOfFrame => {
               if !self.black_screen {
                  self.dirty = true;
               }
            },
            Event::PutPixels(vnc_rect, ref pixels) => {
               self.black_screen = false;
               for x in vnc_rect.left .. (vnc_rect.left+vnc_rect.width) {
                  for y in vnc_rect.top .. (vnc_rect.top+vnc_rect.height) {
                     let i = x - vnc_rect.left;
                     let j = y - vnc_rect.top;
                     let left = 4*(j * vnc_rect.width + i) as usize;
                     if pixels[left]>20 { self.black_screen = false }
                     if pixels[left+1]>20 { self.black_screen = false }
                     if pixels[left+2]>20 { self.black_screen = false }
                  }
               }

               if !self.black_screen {
               for x in vnc_rect.left .. min((self.offset_width+width) as u16, (vnc_rect.left+vnc_rect.width)) {
                  for y in vnc_rect.top .. min((self.offset_height+height) as u16, (vnc_rect.top+vnc_rect.height)) {
                     let i = x - vnc_rect.left;
                     let j = y - vnc_rect.top;
                     let left = 4*(j * vnc_rect.width + i) as usize;
                     if (x as usize)>=self.offset_width && (y as usize)>=self.offset_height {
                        let x = (x as usize) - self.offset_width;
                        let y = (y as usize) - self.offset_height;
                        self.state.screen[3*(y*width + x)] = pixels[left+2];
                        self.state.screen[3*(y*width + x)+1] = pixels[left+1];
                        self.state.screen[3*(y*width + x)+2] = pixels[left];
                     }
                  }
               }}

            },
            _ => {}
         }
      }
   }
}

impl Gym {
   pub fn parse_args(&mut self) -> () {
      let mut prev = "".to_string();
      for argument in std::env::args() {
         if prev == "--fps" {
            self.fps = argument.parse::<u32>().unwrap();
            prev = "".to_string();
         } else if prev == "--game" {
            self.env_id = argument;
            prev = "".to_string();
         } else if prev == "--parallel" {
            self.max_parallel = argument.parse::<u32>().unwrap();
            prev = "".to_string();
         } else if prev == "--duration" {
            self.duration = argument.parse::<u64>().unwrap();
            prev = "".to_string();
         } else if prev == "--record" {
            self.record_dst = argument;
            prev = "".to_string();
         } else {
            prev = argument;
         }
      }
   }
   pub fn set_fps(&mut self, fps: u32) -> () { self.fps = fps }
   pub fn set_game(&mut self, game: String) -> () { self.env_id = game.to_string() }
   pub fn set_max_parallel(&mut self, par: u32) -> () { self.max_parallel = par }
   pub fn set_duration(&mut self, num: u64) -> () { self.duration = num }
   pub fn set_record(&mut self, dst: String) -> () { self.record_dst = dst.to_string() }
   pub fn new() -> Gym {
      Gym {
         fps: 10,
         env_id: "gym-core.AirRaid-v0".to_string(),
         max_parallel: 1,
         duration: 180,
         record_dst: "video.mpg".to_string()
      }
   }
   pub fn remote_prep_container(&mut self, pi: u32) -> () {
      let base_vnc = 5900 + pi;
      let base_rec = 15900 + pi;
      if self.env_id.starts_with("flashgames.") {
         Command::new("docker")
         .arg("run")
         .arg("-p").arg( format!("5900:{}", base_vnc) )
         .arg("-p").arg( format!("15900:{}", base_rec) )
         .arg("--ipc")
         .arg("host")
         .arg("--cap-add")
         .arg("SYS_ADMIN")
         .arg("--privileged")
         .arg("quay.io/openai/universe.flashgames:0.20.25")
         .spawn();
      } else {
         Command::new("docker")
         .arg("run")
         .arg("-p").arg( format!("5900:{}", base_vnc) )
         .arg("-p").arg( format!("15900:{}", base_rec) )
         .arg("quay.io/openai/universe.gym-core:0.20.6")
         .spawn();
      }
      print!("spawned docker process at {} / {}", base_vnc, base_rec);

      for pi in 0..self.max_parallel {
         let mut ok = false;
         for _ in 0..15 { //wait up to 15 seconds for dockers to boot
            if !ok {
               let rec_port = 15900 + pi;

               println!("Polling docker at port {} for connectivity.", rec_port);
               let one_second = time::Duration::from_millis(1000);
               thread::sleep(one_second);

               let ns1: SocketAddr = format!("127.0.0.1:{}", rec_port).parse().expect("Unable to parse socket address");
               let mut s1 = TcpStream::connect(ns1);

               match s1 {
                  Ok(_) => { ok=true; }
                  _ => { println!("No connectivity to {}", rec_port); }
               }
            }
         }
         if !ok { panic!("Unable to confirm connectivity to docker #{}", pi) }
         else { println!("Confirmed connectivity to docker #{}", pi); }
      }

      let ten_seconds = time::Duration::from_millis(10000);
      thread::sleep(ten_seconds);

   }
   pub fn remote_prep_recorder(&mut self, pi: u32) {
      for entry in glob("mov_out/*.png").expect("Failed to read glob pattern") {
         match entry {
            Ok(path) => {
               fs::remove_file(path);
            }
            Err(e) => {}
         }
      }
      for entry in glob("*.mpg").expect("Failed to read glob pattern") {
         match entry {
            Ok(path) => {
               fs::remove_file(path);
            }
            Err(e) => {}
         }
      }
      fs::create_dir("mov_out/");
   }
   pub fn start_remote(&mut self, pi: u32) -> GymRemote {
      self.remote_prep_container(pi);
      self.remote_prep_recorder(pi);
      let mut screen_width = ATARI_WIDTH;
      let mut screen_height = ATARI_HEIGHT;
      let mut offset_width = 0;
      let mut offset_height = 0;
      let mut mode = "atari";
      if self.env_id.starts_with("flashgames.") {
         let data = Json::from_str(include_str!("../flashgames.json")).unwrap();
         let obj = data.as_object().unwrap()
                   .get(&self.env_id[..]).unwrap()
                   .as_object().unwrap();
         screen_width = obj.get("width").unwrap().as_u64().unwrap() as u32;
         screen_height = obj.get("height").unwrap().as_u64().unwrap() as u32;
         offset_width = 20;
         offset_height = 84;
         mode = "flash";
      }
      let mut action_space: std::vec::Vec<usize> = Vec::new();
      if mode=="flash" {
         //mouse input
         action_space.push(screen_width as usize);
         action_space.push(screen_height as usize);
         action_space.push(2);

         //keyboard input
         action_space.push(2); //left
         action_space.push(2); //right
         action_space.push(2); //up
         action_space.push(2); //down
         action_space.push(2); //spacebar
      } else {
         action_space.push(2); //joystick left
         action_space.push(2); //joystick right
         action_space.push(2); //joystick up
         action_space.push(2); //joystick down
         action_space.push(2); //red button
      }
      let r = GymRemote {
         black_screen: false,
         dirty: true,
         id: pi,
         fps: self.fps,
         env_id: format!("{: <99}", self.env_id).clone(),
         duration: self.duration,
         record_dst: self.record_dst.clone(),
         mode: mode.to_string(),
         vnc: None,
         rewarder: None,
         state: GymState {
            screen: vec![0; (screen_width * screen_height * 3) as usize]
         },
         shape: GymShape {
            action_space: action_space,
            observation_space: vec![screen_width as usize, screen_height as usize, 3 as usize],
            reward_max : f64::INFINITY,
            reward_min : f64::NEG_INFINITY
         },
         offset_width: offset_width,
         offset_height: offset_height,
         time: Instant::now(),
         frame: 0
      };
      return r;
   }
   pub fn remote_sanitize(&mut self) -> () {
      Command::new("sh").arg("-c").arg("docker kill $(docker ps -q); true").spawn();
   }
   pub fn start<F,T: GymMember>(&mut self, start_agent: F) -> ()
   where F: Fn() -> T + Send + Sync + 'static
   {
      let start_agent = Arc::new(start_agent);
      let mut threads = Vec::new();
      self.remote_sanitize();
      for pi in 0..self.max_parallel {
         let start_agent = start_agent.clone();
         let mut remote = self.start_remote(pi);
         threads.push(std::thread::spawn(move || {
            let agent = start_agent();
            remote.start(agent);
         }));
      }

      for t in threads {
         t.join();
      }

      self.remote_sanitize();
   }
}