use std::io::{self, Read, Write};
use std::net::{Shutdown, TcpStream};

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

use metainfo::Hash;
use tracker::Peer;

#[derive(Debug)]
pub enum Packet {
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have {
        index: u32,
    },
    Bitfield { bitfield: Vec<u8>, },
    Request { index: u32, begin: u32, length: u32, },
    Piece { index: u32, begin: u32, block: Vec<u8>, },
    Cancel { index: u32, begin: u32, length: u32, }
}

#[derive(Debug)]
pub struct PeerConnection {
    sock: TcpStream,
}

impl PeerConnection {
    pub fn open_connection(peer: Peer, own_info_hash: Hash, own_peer_id: Hash) -> io::Result<(PeerConnection, PeerConnection, Hash)> {
        let mut sock = TcpStream::connect((peer.addr, peer.port))?;

        // send handshake
        sock.write_u8(19)?;
        sock.write(b"BitTorrent protocol")?;
        sock.write_u64::<BigEndian>(0)?;
        sock.write(&own_info_hash)?;
        sock.write(&own_peer_id)?;
        sock.flush()?;

        // receive handshake
        let mut buf = [0u8; 68];
        sock.read_exact(&mut buf)?;

        let peer_info_hash = Hash::from_slice(&buf[28..48]);
        let peer_id = Hash::from_slice(&buf[48..68]);

        if buf[0] != 19 || &buf[1..20] != b"BitTorrent protocol" || own_info_hash != peer_info_hash {
            return Err(io::Error::new(io::ErrorKind::Other, "invalid protocol"))
        }

        let mut conn = PeerConnection { sock: sock };

        Ok((conn.try_clone()?, conn, peer_id))
    }

    pub fn read_packet(&mut self) -> io::Result<Packet> {
        let len = self.sock.read_u32::<BigEndian>()?;
        let id = self.sock.read_u8()?;

        Ok(match id {
            0 => Packet::Choke,
            1 => Packet::Unchoke,
            2 => Packet::Interested,
            3 => Packet::NotInterested,
            4 => Packet::Have {
                index: self.sock.read_u32::<BigEndian>()?,
            },
            5 => {
                let size = len as usize - 1;
                let mut bitfield = Vec::with_capacity(size);
                unsafe {
                    bitfield.set_len(size);
                }
                self.sock.read_exact(&mut bitfield)?;
                Packet::Bitfield {
                    bitfield: bitfield,
                }
            }
            6 => Packet::Request {
                index: self.sock.read_u32::<BigEndian>()?,
                begin: self.sock.read_u32::<BigEndian>()?,
                length: self.sock.read_u32::<BigEndian>()?,
            },
            7 => {
                let size = len as usize - 9;
                let index = self.sock.read_u32::<BigEndian>()?;
                let begin = self.sock.read_u32::<BigEndian>()?;
                let mut block = Vec::with_capacity(size);
                unsafe {
                    block.set_len(size);
                }
                self.sock.read_exact(&mut block)?;

                Packet::Piece {
                    index: index,
                    begin: begin,
                    block: block,
                }
            }
            8 => Packet::Cancel {
                index: self.sock.read_u32::<BigEndian>()?,
                begin: self.sock.read_u32::<BigEndian>()?,
                length: self.sock.read_u32::<BigEndian>()?,
            },
            _ => return Err(io::Error::new(io::ErrorKind::Other, "invalid packet")),
        })
    }

    pub fn send_keepalive(&mut self) -> io::Result<()> {
        self.sock.write(b"")?;
        Ok(())
    }

    pub fn send_choke(&mut self) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(1)?;
        self.sock.write_u8(0)?;
        Ok(())
    }

    pub fn send_unchoke(&mut self) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(1)?;
        self.sock.write_u8(1)?;
        Ok(())
    }

    pub fn send_interested(&mut self) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(1)?;
        self.sock.write_u8(2)?;
        Ok(())
    }

    pub fn send_not_interested(&mut self) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(1)?;
        self.sock.write_u8(3)?;
        Ok(())
    }

    pub fn send_have(&mut self, piece: u32) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(5)?;
        self.sock.write_u8(4)?;
        self.sock.write_u32::<BigEndian>(piece)?;
        Ok(())
    }

    pub fn send_request(&mut self, piece: u32, begin: u32, length: u32) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(13)?;
        self.sock.write_u8(6)?;
        self.sock.write_u32::<BigEndian>(piece)?;
        self.sock.write_u32::<BigEndian>(begin)?;
        self.sock.write_u32::<BigEndian>(length)?;
        Ok(())
    }

    pub fn send_piece(&mut self, piece: u32, begin: u32, data: &[u8]) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(9 + data.len() as u32)?;
        self.sock.write_u8(7)?;
        self.sock.write_u32::<BigEndian>(piece)?;
        self.sock.write_u32::<BigEndian>(begin)?;
        self.sock.write(data)?;
        Ok(())
    }

    pub fn send_cancel(&mut self, piece: u32, begin: u32, length: u32) -> io::Result<()> {
        self.sock.write_u32::<BigEndian>(13)?;
        self.sock.write_u8(8)?;
        self.sock.write_u32::<BigEndian>(piece)?;
        self.sock.write_u32::<BigEndian>(begin)?;
        self.sock.write_u32::<BigEndian>(length)?;
        Ok(())
    }

    pub fn try_clone(&mut self) -> io::Result<PeerConnection> {
        Ok(PeerConnection {
            sock: self.sock.try_clone()?,
        })
    }

    pub fn shutdown(&mut self) {
        let _ = self.sock.shutdown(Shutdown::Both);
    }
}