use std::cell::RefCell;
use std::rc::Rc;

use js_sys::{ArrayBuffer, Uint8Array};
use wasm_bindgen::prelude::*;
use wasm_bindgen::JsCast;
use web_sys::{
    console,
    window,
    BinaryType,
    Document,
    Element,
    ErrorEvent,
    Event,
    HtmlElement,
    HtmlTextAreaElement,
    KeyboardEvent,
    MessageEvent,
    WebSocket,
};

use frontend_protocol::{
    UserChatMessage,
    DekuBytes,
    UserChatCompletionRequest,
    UserRequest,
    UserResponse,
};

thread_local! {
    static APP_HANDLE: RefCell<Option<Rc<AppState>>> = RefCell::new(None);
}

struct ChatState {
    messages: Vec<UserChatMessage>,
    message_nodes: Vec<MessageNode>,
    is_receiving: bool,
    active_assistant_index: Option<usize>,
    compose_role: ComposeRole,
    generation_enabled: bool,
}

struct MessageNode {
    wrapper: Element,
    content: Element,
    status: Option<Element>,
}

#[derive(Copy, Clone)]
enum ChatStatus {
    Pending,
    Cancelled,
    Hidden,
}

#[derive(Copy, Clone, PartialEq)]
enum ComposeRole {
    System,
    Assistant,
    User,
}

impl ComposeRole {
    fn as_str(self) -> &'static str {
        match self {
            ComposeRole::System => "system",
            ComposeRole::Assistant => "assistant",
            ComposeRole::User => "user",
        }
    }

    fn next(self) -> Self {
        match self {
            ComposeRole::System => ComposeRole::Assistant,
            ComposeRole::Assistant => ComposeRole::User,
            ComposeRole::User => ComposeRole::System,
        }
    }
}

struct AppState {
    ws: WebSocket,
    document: Document,
    messages_container: Element,
    input: HtmlTextAreaElement,
    status_label: Element,
    state: RefCell<ChatState>,
}

fn append_message(
    document: &Document,
    container: &Element,
    role: &str,
    content: &str,
    status: Option<ChatStatus>,
) -> Result<MessageNode, JsValue> {
    let wrapper = document.create_element("div")?;
    wrapper.set_class_name(&format!("chat-message chat-message--{}", role));

    let role_el = document.create_element("div")?;
    role_el.set_class_name("chat-message__role");
    role_el.set_text_content(Some(role));

    let content_el = document.create_element("div")?;
    content_el.set_class_name("chat-message__content");
    content_el.set_text_content(Some(content));

    wrapper.append_child(&role_el)?;
    wrapper.append_child(&content_el)?;

    let status_el = if let Some(status) = status {
        let status_el = document.create_element("div")?;
        apply_status(&status_el, status);
        wrapper.append_child(&status_el)?;
        Some(status_el)
    } else {
        None
    };

    container.append_child(&wrapper)?;

    Ok(MessageNode {
        wrapper,
        content: content_el,
        status: status_el,
    })
}

fn set_message_content(node: &Element, content: &str) {
    node.set_text_content(Some(content));
}

fn apply_status(node: &Element, status: ChatStatus) {
    match status {
        ChatStatus::Pending => {
            node.set_class_name("chat-message__status chat-message__status--pending");
            node.set_text_content(Some("..."));
        }
        ChatStatus::Cancelled => {
            node.set_class_name("chat-message__status chat-message__status--cancelled");
            node.set_text_content(Some("[canceled]"));
        }
        ChatStatus::Hidden => {
            node.set_class_name("chat-message__status chat-message__status--hidden");
            node.set_text_content(Some(""));
        }
    }
}

fn update_status_label(label: &Element, role: ComposeRole, generation_enabled: bool) {
    let generation = if generation_enabled { "on" } else { "off" };
    let text = format!("Role: {} | Generation: {}", role.as_str(), generation);
    label.set_text_content(Some(&text));
}

fn scroll_to_bottom(container: &Element) {
    if let Some(element) = container.dyn_ref::<HtmlElement>() {
        let height = element.scroll_height();
        element.set_scroll_top(height);
    }
}

#[wasm_bindgen]
pub fn init_chat() -> Result<(), JsValue> {
    let window = window().ok_or_else(|| JsValue::from_str("Missing window"))?;
    let document = window.document().ok_or_else(|| JsValue::from_str("Missing document"))?;
    let location = window.location();
    let host = location.host()?;
    let protocol = location.protocol()?;
    let scheme = if protocol == "https:" { "wss" } else { "ws" };
    let ws_url = format!("{scheme}://{host}/chat");

    let messages_container = document
        .get_element_by_id("chat-messages")
        .ok_or_else(|| JsValue::from_str("Missing chat-messages element"))?;
    let input_el: HtmlTextAreaElement = document
        .get_element_by_id("chat-input")
        .ok_or_else(|| JsValue::from_str("Missing chat-input element"))?
        .dyn_into()?;
    let status_label = document
        .get_element_by_id("chat-status")
        .ok_or_else(|| JsValue::from_str("Missing chat-status element"))?;

    let ws = WebSocket::new(&ws_url)?;
    ws.set_binary_type(BinaryType::Arraybuffer);
    console::log_1(&format!("[ws] connecting to {ws_url}").into());

    let app = Rc::new(AppState {
        ws,
        document,
        messages_container,
        input: input_el,
        status_label,
        state: RefCell::new(ChatState {
            messages: Vec::new(),
            message_nodes: Vec::new(),
            is_receiving: false,
            active_assistant_index: None,
            compose_role: ComposeRole::User,
            generation_enabled: true,
        }),
    });

    {
        let state = app.state.borrow();
        update_status_label(&app.status_label, state.compose_role, state.generation_enabled);
    }

    let onopen = Closure::<dyn FnMut(Event)>::wrap(Box::new(move |_| {
        console::log_1(&"[ws] connected".into());
    }));
    app.ws.set_onopen(Some(onopen.as_ref().unchecked_ref()));
    onopen.forget();

    let app_for_keydown = app.clone();
    let onkeydown = Closure::<dyn FnMut(KeyboardEvent)>::wrap(Box::new(move |event: KeyboardEvent| {
        if event.ctrl_key() {
            let key = event.key();
            if key == "c" || key == "C" {
                let state = app_for_keydown.state.borrow();
                if state.is_receiving {
                    event.prevent_default();
                    if app_for_keydown.ws.ready_state() != WebSocket::OPEN {
                        console::error_1(&"[ws] socket is not open".into());
                        return;
                    }

                    let request = UserRequest::ChatCompletionCancellation;
                    match request.to_bytes() {
                        Ok(bytes) => {
                            if let Err(err) = app_for_keydown.ws.send_with_u8_array(&bytes) {
                                console::error_1(&format!("[ws] cancel send error: {:?}", err).into());
                            }
                        }
                        Err(err) => {
                            console::error_1(&format!("[ws] cancel encode error: {err:#}").into());
                        }
                    }
                }
                return;
            }

            if key == "1" {
                event.prevent_default();
                let mut state = app_for_keydown.state.borrow_mut();
                state.compose_role = state.compose_role.next();
                update_status_label(
                    &app_for_keydown.status_label,
                    state.compose_role,
                    state.generation_enabled,
                );
                return;
            }

            if key == "4" {
                event.prevent_default();
                let mut state = app_for_keydown.state.borrow_mut();
                state.generation_enabled = !state.generation_enabled;
                update_status_label(
                    &app_for_keydown.status_label,
                    state.compose_role,
                    state.generation_enabled,
                );
                return;
            }

            if key == "d" || key == "D" {
                event.prevent_default();
                let mut state = app_for_keydown.state.borrow_mut();
                if state.is_receiving {
                    return;
                }
                let last_message = match state.messages.pop() {
                    Some(message) => message,
                    None => return,
                };
                if let Some(node) = state.message_nodes.pop() {
                    if let Some(parent) = node.wrapper.parent_node() {
                        let _ = parent.remove_child(&node.wrapper);
                    }
                }
                app_for_keydown.input.set_value(&last_message.content);
                scroll_to_bottom(&app_for_keydown.messages_container);
                return;
            }
        }

        if event.key() != "Enter" || event.shift_key() {
            return;
        }
        event.prevent_default();

        let mut state = app_for_keydown.state.borrow_mut();
        if state.is_receiving {
            return;
        }

        let raw = app_for_keydown.input.value();
        let trimmed = raw.trim();
        if trimmed.is_empty() {
            return;
        }

        let role = state.compose_role;
        let user_content = trimmed.to_string();
        app_for_keydown.input.set_value("");

        let user_node = match append_message(
            &app_for_keydown.document,
            &app_for_keydown.messages_container,
            role.as_str(),
            &user_content,
            None,
        ) {
            Ok(node) => node,
            Err(err) => {
                console::error_1(&format!("[ui] failed to append user message: {:?}", err).into());
                return;
            }
        };

        state.messages.push(UserChatMessage {
            role: role.as_str().to_string(),
            content: user_content,
        });
        state.message_nodes.push(user_node);
        scroll_to_bottom(&app_for_keydown.messages_container);

        if role != ComposeRole::User || !state.generation_enabled {
            return;
        }

        if app_for_keydown.ws.ready_state() != WebSocket::OPEN {
            console::error_1(&"[ws] socket is not open".into());
            return;
        }

        let history = state.messages.clone();
        let request = UserRequest::ChatCompletion(UserChatCompletionRequest::new(history));

        let bytes = match request.to_bytes() {
            Ok(bytes) => bytes,
            Err(err) => {
                console::error_1(&format!("[ws] encode error: {err:#}").into());
                return;
            }
        };

        if let Err(err) = app_for_keydown.ws.send_with_u8_array(&bytes) {
            console::error_1(&format!("[ws] send error: {:?}", err).into());
            return;
        }

        let assistant_node = match append_message(
            &app_for_keydown.document,
            &app_for_keydown.messages_container,
            "assistant",
            "",
            Some(ChatStatus::Pending),
        ) {
            Ok(node) => node,
            Err(err) => {
                console::error_1(&format!("[ui] failed to append assistant message: {:?}", err).into());
                return;
            }
        };

        state.messages.push(UserChatMessage {
            role: "assistant".to_string(),
            content: String::new(),
        });
        state.message_nodes.push(assistant_node);
        state.active_assistant_index = Some(state.messages.len() - 1);
        state.is_receiving = true;
        scroll_to_bottom(&app_for_keydown.messages_container);
    }));
    app.input.set_onkeydown(Some(onkeydown.as_ref().unchecked_ref()));
    onkeydown.forget();

    let app_for_message = app.clone();
    let onmessage = Closure::<dyn FnMut(MessageEvent)>::wrap(Box::new(move |event: MessageEvent| {
        let data = event.data();
        if let Some(text) = data.as_string() {
            console::log_1(&format!("[ws] unexpected text frame: {text}").into());
            return;
        }
        if !data.is_instance_of::<ArrayBuffer>() {
            console::log_1(&"[ws] unexpected non-binary frame".into());
            return;
        }

        let data = Uint8Array::new(&data);
        let bytes = data.to_vec();
        let response = match UserResponse::from_bytes(&bytes) {
            Ok(response) => response,
            Err(err) => {
                console::error_1(&format!("[ws] decode error: {err:#}").into());
                return;
            }
        };

        let mut state = app_for_message.state.borrow_mut();
        let assistant_index = match state.active_assistant_index {
            Some(index) => index,
            None => {
                console::log_1(&"[ws] missing assistant index".into());
                return;
            }
        };

        match response {
            UserResponse::ChatCompletion(completion) => {
                if let Some(message) = state.messages.get_mut(assistant_index) {
                    message.content.push_str(&completion.piece);
                }
                if let Some(node) = state.message_nodes.get(assistant_index) {
                    if let Some(message) = state.messages.get(assistant_index) {
                        set_message_content(&node.content, &message.content);
                    }
                }
                scroll_to_bottom(&app_for_message.messages_container);
            }
            UserResponse::ChatCompletionEnd => {
                state.is_receiving = false;
                if let Some(node) = state
                    .active_assistant_index
                    .and_then(|index| state.message_nodes.get(index))
                {
                    if let Some(status_node) = node.status.as_ref() {
                        apply_status(status_node, ChatStatus::Hidden);
                    }
                }
                state.active_assistant_index = None;
                let _ = app_for_message.input.focus();
            }
            UserResponse::ChatCompletionCancellation => {
                state.is_receiving = false;
                if let Some(node) = state
                    .active_assistant_index
                    .and_then(|index| state.message_nodes.get(index))
                {
                    if let Some(status_node) = node.status.as_ref() {
                        apply_status(status_node, ChatStatus::Cancelled);
                    }
                }
                state.active_assistant_index = None;
                let _ = app_for_message.input.focus();
            }
        }
    }));
    app.ws.set_onmessage(Some(onmessage.as_ref().unchecked_ref()));
    onmessage.forget();

    let onerror = Closure::<dyn FnMut(ErrorEvent)>::wrap(Box::new(move |event: ErrorEvent| {
        console::error_1(&format!("[ws] error: {}", event.message()).into());
    }));
    app.ws.set_onerror(Some(onerror.as_ref().unchecked_ref()));
    onerror.forget();

    APP_HANDLE.with(|slot| *slot.borrow_mut() = Some(app.clone()));
    let _ = app.input.focus();
    Ok(())
}
