use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use anyhow::{anyhow, Context, Result};
use futures::future::BoxFuture;
use russh::ChannelMsg;
use tokio::io::AsyncWriteExt;
use tokio::sync::{broadcast, mpsc, RwLock};

use super::super::super::models::config::SshConfig;
use super::super::super::models::session::{PtyOptions, SessionEvent, SessionId};
use super::super::super::models::types::PtyOpenRequest;
use super::super::super::traits::PtyService;
use super::connect::{connect_and_auth, ClientHandler};
use super::exec::disconnect_session;

/// Default PTY terminal service implementation.
///
/// This type is primarily for GUI, Tauri, or other scenarios requiring a "real terminal session".
/// For single command execution, prefer using `RusshTool` instead.
#[derive(Clone, Debug)]
pub struct RusshTerminalService {
    state: Arc<ServiceState>,
}

#[derive(Debug, Default)]
struct ServiceState {
    next_session_id: AtomicU64,
    sessions: RwLock<HashMap<SessionId, SessionHandle>>,
}

#[derive(Clone, Debug)]
struct SessionHandle {
    command_tx: mpsc::Sender<SessionCommand>,
    event_tx: broadcast::Sender<SessionEvent>,
}

#[derive(Debug)]
enum SessionCommand {
    Input(Vec<u8>),
    Resize {
        cols: u32,
        rows: u32,
        pix_width: u32,
        pix_height: u32,
    },
    Close,
}

impl Default for RusshTerminalService {
    fn default() -> Self {
        Self {
            state: Arc::new(ServiceState {
                next_session_id: AtomicU64::new(1),
                sessions: RwLock::new(HashMap::new()),
            }),
        }
    }
}

impl RusshTerminalService {
    /// Create a new terminal service instance.
    pub fn new() -> Self {
        Self::default()
    }

    fn next_session_id(&self) -> SessionId {
        SessionId::new(self.state.next_session_id.fetch_add(1, Ordering::Relaxed))
    }

    async fn session_handle(&self, session_id: SessionId) -> Result<SessionHandle> {
        let guard = self.state.sessions.read().await;
        guard
            .get(&session_id)
            .cloned()
            .ok_or_else(|| anyhow!("session {} not found", session_id.as_u64()))
    }
}

impl PtyService for RusshTerminalService {
    fn open(&self, cfg: SshConfig, req: PtyOpenRequest) -> BoxFuture<'static, Result<SessionId>> {
        let service = self.clone();
        Box::pin(async move {
            let command = req.command.clone();
            let options: PtyOptions = req.into();
            let session = connect_and_auth(cfg).await?;
            let channel = session
                .channel_open_session()
                .await
                .with_context(|| "open session channel failed")?;

            channel
                .request_pty(
                    true,
                    &options.term,
                    options.cols,
                    options.rows,
                    options.pix_width,
                    options.pix_height,
                    &[],
                )
                .await
                .with_context(|| "request pty failed")?;
            if let Some(command) = command {
                channel
                    .exec(true, command.as_str())
                    .await
                    .with_context(|| "exec in pty failed")?;
            } else {
                channel
                    .request_shell(true)
                    .await
                    .with_context(|| "request shell failed")?;
            }

            let session_id = service.next_session_id();
            let (command_tx, command_rx) = mpsc::channel(128);
            let (event_tx, _) = broadcast::channel(4096);

            {
                let mut guard = service.state.sessions.write().await;
                guard.insert(
                    session_id,
                    SessionHandle {
                        command_tx,
                        event_tx: event_tx.clone(),
                    },
                );
            }

            tokio::spawn(run_terminal_session(
                service.state.clone(),
                session_id,
                session,
                channel,
                command_rx,
                event_tx,
            ));

            Ok(session_id)
        })
    }

    fn write_input(&self, session_id: SessionId, data: Vec<u8>) -> BoxFuture<'static, Result<()>> {
        let this = self.clone();
        Box::pin(async move {
            let handle = this.session_handle(session_id).await?;
            handle
                .command_tx
                .send(SessionCommand::Input(data))
                .await
                .map_err(|_| anyhow!("session {} input channel closed", session_id.as_u64()))?;
            Ok(())
        })
    }

    fn resize(
        &self,
        session_id: SessionId,
        cols: u32,
        rows: u32,
        pix_width: u32,
        pix_height: u32,
    ) -> BoxFuture<'static, Result<()>> {
        let this = self.clone();
        Box::pin(async move {
            let handle = this.session_handle(session_id).await?;
            handle
                .command_tx
                .send(SessionCommand::Resize {
                    cols,
                    rows,
                    pix_width,
                    pix_height,
                })
                .await
                .map_err(|_| anyhow!("session {} resize channel closed", session_id.as_u64()))?;
            Ok(())
        })
    }

    fn close(&self, session_id: SessionId) -> BoxFuture<'static, Result<()>> {
        let this = self.clone();
        Box::pin(async move {
            let handle = this.session_handle(session_id).await?;
            handle
                .command_tx
                .send(SessionCommand::Close)
                .await
                .map_err(|_| anyhow!("session {} close channel closed", session_id.as_u64()))?;
            Ok(())
        })
    }

    fn subscribe(
        &self,
        session_id: SessionId,
    ) -> BoxFuture<'static, Result<broadcast::Receiver<SessionEvent>>> {
        let this = self.clone();
        Box::pin(async move {
            let handle = this.session_handle(session_id).await?;
            Ok(handle.event_tx.subscribe())
        })
    }
}

async fn run_terminal_session(
    state: Arc<ServiceState>,
    session_id: SessionId,
    session: russh::client::Handle<ClientHandler>,
    mut channel: russh::Channel<russh::client::Msg>,
    mut command_rx: mpsc::Receiver<SessionCommand>,
    event_tx: broadcast::Sender<SessionEvent>,
) {
    let mut writer = channel.make_writer();
    let mut exit_status = None;

    loop {
        tokio::select! {
            maybe_command = command_rx.recv() => {
                match maybe_command {
                    Some(SessionCommand::Input(data)) => {
                        let write_result = writer.write_all(&data).await;
                        let flush_result = if write_result.is_ok() {
                            writer.flush().await
                        } else {
                            Ok(())
                        };

                        if let Err(err) = write_result.and(flush_result) {
                            let _ = event_tx.send(SessionEvent::Error(format!("write to remote pty failed: {err:#}")));
                            break;
                        }
                    }
                    Some(SessionCommand::Resize { cols, rows, pix_width, pix_height }) => {
                        if let Err(err) = channel.window_change(cols, rows, pix_width, pix_height).await {
                            let _ = event_tx.send(SessionEvent::Error(format!("resize remote pty failed: {err:#}")));
                            break;
                        }
                    }
                    Some(SessionCommand::Close) | None => {
                        channel.eof().await.ok();
                        channel.close().await.ok();
                        break;
                    }
                }
            }
            maybe_msg = channel.wait() => {
                match maybe_msg {
                    Some(ChannelMsg::Data { data }) => {
                        let _ = event_tx.send(SessionEvent::Stdout(data.to_vec()));
                    }
                    Some(ChannelMsg::ExtendedData { data, .. }) => {
                        let _ = event_tx.send(SessionEvent::Stderr(data.to_vec()));
                    }
                    Some(ChannelMsg::ExitStatus { exit_status: code }) => {
                        exit_status = Some(code);
                    }
                    Some(ChannelMsg::Close) | None => {
                        break;
                    }
                    _ => {}
                }
            }
        }
    }

    channel.close().await.ok();
    disconnect_session(&session).await.ok();

    let _ = event_tx.send(SessionEvent::Closed { exit_status });
    state.sessions.write().await.remove(&session_id);
}
