use std::error::Error;
use std::fmt::{Display, Formatter};
use std::io;
use std::io::{Chain, Cursor, ErrorKind, Read, Write};
use std::string::FromUtf8Error;
use crate::frame::{Frame, Message, Opcode};
use crate::frame::frame_ops::{AsFrame, FromFrame};
use crate::frame::read::FrameReadError;
use crate::frame::wire::{FrameReadable, FrameWritable};
use crate::frame::write::FrameWriteError;
use crate::random::masking_key;
use crate::splittable::Splittable;

/// An upgraded TCP stream, used to send WebSocket messages.
pub struct WebsocketStream<R: Read, W: Write> {
    enable_masking: bool,
    read: R,
    write: W,
    is_in_fragment: bool,
    fragment_buf: Vec<u8>,
    fragment_opcode: Opcode,
    closing: bool,
    closed: bool,
    close_code: Option<u16>,
    close_message: Option<String>
}

impl<S: Read + Write + Splittable> WebsocketStream<Chain<Cursor<Vec<u8>>, S>, S> {
    pub(crate) fn create(mask: bool, stream: S, extra_bytes: Vec<u8>) -> Self {
        let extra_bytes_cur = Cursor::new(extra_bytes);
        Self {
            enable_masking: mask,
            read: extra_bytes_cur.chain(stream.try_split().unwrap()),
            write: stream,
            is_in_fragment: false,
            fragment_buf: vec![],
            fragment_opcode: Opcode::Continuation,
            closed: false,
            closing: false,
            close_code: None,
            close_message: None
        }
    }
}

impl<R: Read, W: Write> FrameWritable for WebsocketStream<R, W> {
    type Error = FrameWriteError;

    fn write_frame(&mut self, frame: &Frame) -> Result<(), Self::Error> {
        self.write.write_frame(frame)
    }
}

impl<R: Read, W: Write> FrameReadable for WebsocketStream<R, W> {
    type Error = FrameReadError;

    fn read_frame(&mut self) -> Result<Frame, Self::Error> {
        loop {
            let res = self.read.read_frame();
            match res {
                Ok(f) => return Ok(f),
                Err(e) => {
                    match e {
                        FrameReadError::IoError(e) if matches!(e.kind(), ErrorKind::WouldBlock) => {
                            continue;
                        }
                        _ => return Err(e)
                    }
                }
            }
        }
    }
}

pub trait MessageReadable<M: FromFrame>: FrameReadable {
    type Err;

    fn read_message(&mut self) -> Result<M, Self::Err>;
}
pub trait MessageWritable: FrameWritable {
    type Err;

    fn write_message(&mut self, msg: impl AsFrame) -> Result<(), Self::Err>;
}

#[derive(Debug)]
pub enum MessageReadError {
    /// An IO error occured while reading the message
    IoError(io::Error),
    /// An error occured while reading the frame
    FrameReadError(FrameReadError),
    /// Received a continuation frame while not processing a fragment
    ContinuationOutsideFragment,
    /// The connection was closed
    ConnectionClosed {
        code: Option<u16>,
        message: Option<String>
    },
    /// There was an error writing a response to an internal frame to the stream
    FrameWriteError(FrameWriteError),
    /// There was an error parsing a UTF-8 byte literal
    FromUtf8Error(FromUtf8Error)
}
impl Display for MessageReadError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::IoError(e) => write!(f, "io error: {}", e),
            Self::FrameReadError(e) => write!(f, "frame read error: {}", e),
            Self::ContinuationOutsideFragment => write!(f, "received continuation frame while not in fragment"),
            Self::ConnectionClosed { code, message } => {
                write!(f, "channel closed (code: {:?}, message: {:?})", code, message)
            },
            Self::FrameWriteError(e) => write!(f, "frame write error: {}", e),
            Self::FromUtf8Error(e) => write!(f, "invalid UTF-8: {}", e)
        }
    }
}
impl Error for MessageReadError {}
impl From<io::Error> for MessageReadError {
    fn from(value: io::Error) -> Self {
        Self::IoError(value)
    }
}
impl From<FrameReadError> for MessageReadError {
    fn from(value: FrameReadError) -> Self {
        Self::FrameReadError(value)
    }
}
impl From<FrameWriteError> for MessageReadError {
    fn from(value: FrameWriteError) -> Self {
        Self::FrameWriteError(value)
    }
}
impl From<FromUtf8Error> for MessageReadError {
    fn from(value: FromUtf8Error) -> Self {
        Self::FromUtf8Error(value)
    }
}

impl<R: Read, W: Write> MessageReadable<Message> for WebsocketStream<R, W> {
    type Err = MessageReadError;

    fn read_message(&mut self) -> Result<Message, Self::Err> {
        // loop until we have a full message to send upwards
        loop {
            debug!("reading a frame");
            // read a frame
            let frame = self.read_frame()?;

            debug!("read frame {} bytes, opcode {:?}", frame.payload_len, frame.opcode);

            if frame.opcode == Opcode::Continuation && self.is_in_fragment {
                self.fragment_buf.extend_from_slice(&frame.payload_data);
                if frame.fin {
                    // fragment over
                    let mut frame = frame;
                    frame.payload_data = self.fragment_buf.clone();
                    frame.opcode = self.fragment_opcode.clone();
                    frame.payload_len = frame.payload_data.len() as u64;
                    self.is_in_fragment = false;
                    self.fragment_buf = vec![];
                    self.fragment_opcode = Opcode::Continuation;
                    return Ok(Message::from_frame(&frame));
                }
                // fragment not over, keep loopin'
                continue;
            }

            return match frame.opcode {
                Opcode::Continuation => Err(MessageReadError::ContinuationOutsideFragment),
                Opcode::Text | Opcode::Binary | Opcode::Ping | Opcode::Pong => Ok(Message::from_frame(&frame)),
                Opcode::ConnectionClose => {
                    if !self.closing {
                        // send a response
                        self.write_frame(&Frame {
                            fin: true,
                            rsv1: false,
                            rsv2: false,
                            rsv3: false,
                            opcode: Opcode::ConnectionClose,
                            mask: self.enable_masking,
                            payload_len: 2,
                            masking_key: if self.enable_masking { Some(masking_key()) } else { None },
                            payload_data: vec![0x03, 0xe8],
                        })?;
                    }
                    self.closed = true;
                    if frame.payload_len == 2 {
                        // just has a close reason
                        let close_code = [frame.payload_data[0], frame.payload_data[1]];
                        let close_code = u16::from_be_bytes(close_code);
                        self.close_code = Some(close_code);
                        return Err(MessageReadError::ConnectionClosed {
                            code: Some(close_code),
                            message: None
                        })
                    } else if frame.payload_len > 2 {
                        // close reason and message
                        let close_code = [frame.payload_data[0], frame.payload_data[1]];
                        let close_code = u16::from_be_bytes(close_code);

                        let close_message = &frame.payload_data[2..];
                        let close_message = String::from_utf8(close_message.to_vec())?;
                        self.close_code = Some(close_code);
                        self.close_message = Some(close_message.clone());
                        return Err(MessageReadError::ConnectionClosed {
                            code: Some(close_code),
                            message: Some(close_message)
                        })
                    }
                    Err(MessageReadError::ConnectionClosed {
                        code: None,
                        message: None
                    })
                }
            }
        }
    }
}

#[derive(Debug)]
pub enum MessageWriteError {
    /// An IO error occured while reading the message
    IoError(io::Error),
    /// The connection is closed
    ConnectionClosed {
        code: Option<u16>,
        message: Option<String>
    },
    /// There was an error writing the frame
    FrameWriteError(FrameWriteError),
}
impl Display for MessageWriteError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::IoError(e) => write!(f, "io error: {}", e),
            Self::ConnectionClosed { code, message } => {
                write!(f, "channel closed (code: {:?}, message: {:?})", code, message)
            },
            Self::FrameWriteError(e) => write!(f, "frame write error: {}", e),
        }
    }
}
impl Error for MessageWriteError {}
impl From<io::Error> for MessageWriteError {
    fn from(value: io::Error) -> Self {
        Self::IoError(value)
    }
}
impl From<FrameWriteError> for MessageWriteError {
    fn from(value: FrameWriteError) -> Self {
        Self::FrameWriteError(value)
    }
}

impl<R: Read, W: Write> MessageWritable for WebsocketStream<R, W> {
    type Err = MessageWriteError;

    fn write_message(&mut self, msg: impl AsFrame) -> Result<(), Self::Err> {
        if self.closed {
            return Err(MessageWriteError::ConnectionClosed {
                code: self.close_code,
                message: self.close_message.clone()
            });
        }

        self.write_frame(&msg.as_frame(self.enable_masking))?;

        Ok(())
    }
}

impl<R: Read, W: Write> WebsocketStream<R, W> {
    pub fn close(&mut self, code: Option<u16>, message: Option<String>) -> Result<(), MessageWriteError> {
        self.closing = true;
        self.close_code = code;
        self.close_message = message;

        let mut payload = vec![];
        if let Some(c) = code {
            payload.extend_from_slice(&c.to_be_bytes());
        }
        if let Some(t) = self.close_message.as_ref() {
            payload.extend_from_slice(&t.as_bytes())
        }

        self.write_frame(&Frame {
            fin: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            opcode: Opcode::ConnectionClose,
            mask: self.enable_masking,
            payload_len: payload.len() as u64,
            masking_key: if self.enable_masking { Some(masking_key()) } else { None },
            payload_data: payload,
        })?;

        Ok(())
    }
}