use std::cmp::min;

use ratatui::{
    layout::{Alignment, Constraint, Flex, Layout, Margin, Rect},
    style::{Color, Modifier, Style, Stylize},
    text::{Line, Span, Text},
    widgets::{
        Block, BorderType, Borders, Clear, HighlightSpacing, List, ListItem, Padding, Paragraph,
        Scrollbar, ScrollbarOrientation, ScrollbarState, Wrap,
    },
    Frame,
};
use syntect::highlighting::Theme;
use tui_big_text::{BigText, PixelSize};

use crate::{
    app::{get_file_content, App, AppMode, Message, Notification},
    snippets::{create_highlighted_code, translate_language_name_to_syntect_name},
    storage::list_all_messages,
};
pub const SELECTED_STYLE: Style = Style::new()
    .add_modifier(Modifier::BOLD)
    .fg(Color::LightBlue)
    .bg(Color::DarkGray);

/// helper function to create a centered rect using up certain percentage of the available rect `r`
fn centered_rect(percent_x: u16, percent_y: u16, r: Rect) -> Rect {
    let popup_layout = Layout::vertical([
        Constraint::Percentage((100 - percent_y) / 2),
        Constraint::Percentage(percent_y),
        Constraint::Percentage((100 - percent_y) / 2),
    ])
    .split(r);

    Layout::horizontal([
        Constraint::Percentage((100 - percent_x) / 2),
        Constraint::Percentage(percent_x),
        Constraint::Percentage((100 - percent_x) / 2),
    ])
    .split(popup_layout[1])[1]
}

fn right_aligned_rect(r: Rect, p: u16) -> Rect {
    Layout::horizontal([Constraint::Percentage(100 - p), Constraint::Fill(1)]).split(r)[1]
}

fn left_aligned_rect(r: Rect, p: u16) -> Rect {
    Layout::horizontal([Constraint::Fill(1), Constraint::Percentage(100 - p)]).split(r)[0]
}

fn process_code_blocks<'a>(text: impl Into<String>, width: usize, theme: Theme) -> Vec<Line<'a>> {
    let mut lines = Vec::new();
    let mut in_code_block = false;
    let mut code_buffer = String::new();
    let mut language = String::new();
    let text = text.into();
    let mut nspaces = 0;

    for line in text.lines() {
        if line.trim_start().starts_with("```") {
            if in_code_block {
                // End of code block
                if !code_buffer.is_empty() {
                    let highlighted = if !language.is_empty() {
                        create_highlighted_code(
                            &code_buffer,
                            translate_language_name_to_syntect_name(Some(&language)),
                            &theme,
                        )
                    } else {
                        let wrapped = textwrap::wrap(&code_buffer, width - 3);
                        Text::from(wrapped.join("\n"))
                    };
                    lines.push(
                        Line::from(format!("{}```{}", " ".repeat(nspaces), &language))
                            .style(Style::default().fg(Color::DarkGray)),
                    );
                    lines.extend(highlighted.lines);
                    lines.push(
                        Line::from(format!("{}```", " ".repeat(nspaces)))
                            .style(Style::default().fg(Color::DarkGray)),
                    );
                }
                code_buffer.clear();
                language.clear();
                in_code_block = false;
                // Assume no indentation
                nspaces = 0;
            } else {
                // Start of code block
                nspaces = line.len() - line.trim_start_matches(' ').len();
                language = line.trim_start().trim_start_matches('`').to_string();
                in_code_block = true;
            }
        } else if in_code_block {
            code_buffer.push_str(line);
            code_buffer.push('\n');
        } else {
            // Regular text - wrap it
            let wrapped = textwrap::wrap(line, width - 3);
            lines.extend(wrapped.into_iter().map(|l| Line::from(l.to_string())));
        }
    }
    lines
}

pub fn style_message<'a>(message: Message, width: usize, theme: Theme) -> Vec<Line<'a>> {
    let mut line_vec = Vec::new();
    match message {
        Message::User(text) => {
            line_vec.push(Line::from(Span::raw("USER:").bold().yellow()));
            line_vec.push(Line::from(Span::raw("---").bold().yellow()));
            line_vec.extend(process_code_blocks(text, width, theme));
            line_vec.push(Line::from(Span::raw("").bold().yellow()));
        }
        Message::Assistant(text) => {
            line_vec.push(Line::from(Span::raw("ASSISTANT:").bold().green()));
            line_vec.push(Line::from(Span::raw("---").bold().green()));
            line_vec.extend(process_code_blocks(text, width, theme));
            line_vec.push(Line::from(Span::raw("").bold().green()));
        }
    }
    line_vec
}

fn messages_to_lines(messages: &[Message], width: usize) -> Vec<Line<'_>> {
    let mut line_vec = Vec::new();
    for message in messages {
        match message {
            Message::User(m) => {
                let wrapped_message = textwrap::wrap(m, width - 3);
                line_vec.push(Line::from(Span::raw("USER:").bold().yellow()));
                line_vec.push(Line::from(Span::raw("---").bold().yellow()));
                line_vec.extend(
                    wrapped_message
                        .into_iter()
                        .map(|l| Line::from(Span::raw(l))),
                );
                line_vec.push(Line::from(Span::raw("")));
            }
            Message::Assistant(m) => {
                let wrapped_message = textwrap::wrap(m, width - 3);
                line_vec.push(Line::from(Span::raw("ASSISTANT:").bold().green()));
                line_vec.push(Line::from(Span::raw("---").bold().green()));
                line_vec.extend(
                    wrapped_message
                        .into_iter()
                        .map(|l| Line::from(Span::raw(l))),
                );
                line_vec.push(Line::from(Span::raw("")));
            }
        }
    }
    line_vec
}

fn render_messages(f: &mut Frame, app: &mut App, messages_area: Rect) {
    let scrollbar = Scrollbar::new(ScrollbarOrientation::VerticalRight)
        .begin_symbol(Some("↑"))
        .end_symbol(Some("↓"));
    let messages = if app.is_streaming {
        messages_to_lines(&app.messages, messages_area.width as usize)
    } else {
        app.cached_lines.clone()
    };
    let mut scrollbar_state = ScrollbarState::new(messages.len()).position(app.vertical_scroll);

    let messages_text = Text::from(messages);
    let messages = Paragraph::new(messages_text)
        .scroll((app.vertical_scroll as u16, 0))
        .block(Block::bordered().title(format!("Chat - {}", app.selected_model_name)));

    f.render_widget(messages, messages_area);

    f.render_stateful_widget(
        scrollbar,
        messages_area.inner(Margin {
            // using an inner vertical margin of 1 unit makes the scrollbar inside the block
            vertical: 1,
            horizontal: 0,
        }),
        &mut scrollbar_state,
    );
}

// Fron Ratatui's website
fn center_rect(area: Rect, horizontal: Constraint, vertical: Constraint) -> Rect {
    let [area] = Layout::horizontal([horizontal])
        .flex(Flex::Center)
        .areas(area);
    let [area] = Layout::vertical([vertical]).flex(Flex::Center).areas(area);
    area
}

fn render_init_screen(f: &mut Frame, area: Rect) {
    let big_text = BigText::builder()
        .alignment(Alignment::Center)
        .pixel_size(PixelSize::Full)
        .lines(vec!["AIT".into()])
        .build();
    let centered_area = center_rect(area, Constraint::Length(26), Constraint::Length(8)); // 3 8x8
                                                                                          // characters
    f.render_widget(big_text, centered_area);
}
pub fn render(f: &mut Frame, app: &mut App) {
    let title = format!("AI in the Terminal (AIT v{})", env!("CARGO_PKG_VERSION"));
    let main_block = Block::bordered()
        .title(title)
        .title_alignment(Alignment::Center)
        .border_type(BorderType::Rounded);
    match app.app_mode {
        AppMode::Normal => {
            f.render_widget(main_block.border_style(Style::new().blue()), f.area());
        }
        AppMode::Editing => {
            f.render_widget(main_block.border_style(Style::new().yellow()), f.area());
        }
        _ => {
            f.render_widget(main_block, f.area());
        }
    }

    let input_area_constraint = match app.app_mode {
        AppMode::Editing => Constraint::Min(1),
        _ => Constraint::Length(0),
    };

    let vertical = Layout::vertical([
        Constraint::Length(1),
        Constraint::Min(1),
        input_area_constraint,
    ]);

    let vertical = vertical.margin(1);

    let [help_area, messages_area, input_area] = vertical.areas(f.area());

    match &app.app_mode {
        AppMode::Normal => {
            if !app.messages.is_empty() {
                render_messages(f, app, messages_area);
            } else {
                render_init_screen(f, messages_area);
            }
        }
        AppMode::Editing => {
            render_messages(f, app, messages_area);
            f.render_widget(&app.input_textarea, input_area);
        }
        AppMode::ModelSelection => {
            let block = Block::bordered().title("Select Model");
            let area = centered_rect(40, 50, messages_area);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);
            render_model_list(f, area, app);
        }
        AppMode::SnippetSelection => {
            let block = Block::bordered().title("Select Snippet");
            let area = left_aligned_rect(messages_area, 25);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);
            render_snippet_list(f, area, app);

            let preview_block = Block::bordered().title("Snippet Preview");
            let preview_area = right_aligned_rect(messages_area, 75);
            f.render_widget(Clear, preview_area); //this clears out the background
            f.render_widget(preview_block, preview_area);
            let snippet = app.get_snippet();
            let preview_block_content = Block::new().padding(Padding::uniform(1));
            if let Some(preview_text) = snippet {
                let snippet_text = if let Some(lang) = &preview_text.language {
                    create_highlighted_code(&preview_text.text, lang, &app.theme)
                } else {
                    Text::from(preview_text.text.as_str()).magenta()
                };
                let snippet_paragraph = Paragraph::new(snippet_text).block(preview_block_content);
                f.render_widget(snippet_paragraph, preview_area);
            }
        }
        AppMode::ShowHistory => {
            let block = Block::bordered().title("Select Chat");
            let area = left_aligned_rect(messages_area, 25);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);
            render_chat_history_list(f, area, app);

            let preview_block = Block::bordered().title("Chat Preview");
            let preview_area = right_aligned_rect(messages_area, 75);
            f.render_widget(Clear, preview_area); //this clears out the background
            f.render_widget(preview_block, preview_area);
            let chat_id = app.get_selected_chat_id();
            let preview_text = if let Some(id) = chat_id {
                let text = list_all_messages(*id)
                    .unwrap_or([].to_vec())
                    .into_iter()
                    .map(|m| match m {
                        Message::User(t) => format!("USER: {t}\n"),
                        Message::Assistant(t) => format!("ASSISTANT: {t}\n"),
                    })
                    .collect::<Vec<String>>()
                    .join("\n");
                Some(text)
            } else {
                None
            };
            let preview_block_content = Block::new().padding(Padding::uniform(1));
            if let Some(preview_text) = preview_text {
                let snippet_paragraph = Paragraph::new(Text::from(preview_text.as_str()).magenta())
                    .wrap(Wrap { trim: true })
                    .block(preview_block_content);
                f.render_widget(snippet_paragraph, preview_area);
            }
        }
        AppMode::Help => {
            let block = Block::bordered().title("Help - Use j/k or Up/Down to scroll");
            let area = centered_rect(50, 60, messages_area);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);

            let normal_keys = vec![
                "Press ".into(),
                "Esc/q".bold(),
                " to exit, ".into(),
                "i".bold(),
                " to start editing, ".into(),
                "y".bold(),
                " to copy the last answer (not linux yet), ".into(),
                "m".bold(),
                " to choose model, ".into(),
                "h".bold(),
                " to browse previous conversations, ".into(),
                "s".bold(),
                " to browse code snippets, ".into(),
                "f".bold(),
                " to explore files, ".into(),
                "c".bold(),
                " to view context files, ".into(),
                "n".bold(),
                " to start a new chat, ".into(),
                "r".bold(),
                " to redo last message. ".into(),
                "Scroll with ".into(),
                "j/k or Up/Down".bold(),
                ", ".into(),
                "g".bold(),
                " for top, ".into(),
                "G".bold(),
                " for bottom.".into(),
            ];
            let editing_keys = vec![
                "Press ".into(),
                "Esc".bold(),
                " to stop editing. Press ".into(),
                "CONTROL + S (C-s)".bold(),
                " to submit the message. ".into(),
                "Paste into the text area by pressing ".into(),
                "Ctrl + V".bold(),
            ];
            let model_keys = vec![
                "Press ".into(),
                "Up/Down".bold(),
                " to select model, or press ".into(),
                "Enter".bold(),
                " to select model, which immediately enters 'editing' mode.".into(),
            ];
            let chat_keys = vec![
                "Press ".into(),
                "Up/Down".bold(),
                " to select chat, or press ".into(),
                "d".bold(),
                " to delete the selected chat, or press ".into(),
                "Enter".bold(),
                " to select a chat, and return to 'normal' mode.".into(),
            ];
            let snippet_keys = vec![
                "Press ".into(),
                "Up/Down".bold(),
                " to select snippet, or press ".into(),
                "Enter".bold(),
                " to copy snippet to the clipboard (not linux yet), and return to 'normal' mode."
                    .into(),
            ];
            let file_explorer_keys = vec![
                "Press ".into(),
                "h/j/k/l or arrows".bold(),
                " to navigate directories and files. Press ".into(),
                "Enter".bold(),
                " to add a file to context. Press ".into(),
                "d".bold(),
                " to remove the selected file from context. Press ".into(),
                "Esc/q".bold(),
                " to return to 'normal' mode.".into(),
            ];
            let context_keys = vec![
                "Files added to context will be automatically included in your next message to the LLM. Press ".into(),
                "Esc/q/Enter".bold(),
                " to return to 'normal' mode.".into(),
            ];
            let msg = vec![
                Line::from(Span::raw("Welcome to AI in the Terminal! ").bold()),
                Line::from(""),
                Line::from(vec![
                    "When in ".bold(),
                    "normal".bold().blue(),
                    " mode, you can:".bold(),
                ]),
                Line::from(normal_keys),
                Line::from(""),
                Line::from(vec![
                    "When in ".bold(),
                    "editing".bold().yellow(),
                    " mode, you can:".bold(),
                ]),
                Line::from(editing_keys),
                Line::from(""),
                Line::from(Span::raw("When choosing models, you can:").bold()),
                Line::from(model_keys),
                Line::from(""),
                Line::from(Span::raw("When choosing chats, you can:").bold()),
                Line::from(chat_keys),
                Line::from(""),
                Line::from(Span::raw("When browsing snippets, you can:").bold()),
                Line::from(snippet_keys),
                Line::from(""),
                Line::from(Span::raw("When exploring files, you can:").bold()),
                Line::from(file_explorer_keys),
                Line::from(""),
                Line::from(Span::raw("When viewing context:").bold()),
                Line::from(context_keys),
            ];
            let help_text_block = Block::new().padding(Padding::uniform(1));
            let text = Text::from(msg).patch_style(Style::default());
            let help_message = Paragraph::new(text)
                .scroll((app.help_scroll as u16, 0))
                .block(help_text_block)
                .wrap(Wrap { trim: true });
            f.render_widget(help_message, area);

            // Add scrollbar
            let scrollbar = Scrollbar::new(ScrollbarOrientation::VerticalRight)
                .begin_symbol(Some("↑"))
                .end_symbol(Some("↓"));
            let mut scrollbar_state = ScrollbarState::new(30).position(app.help_scroll);
            f.render_stateful_widget(
                scrollbar,
                area.inner(Margin {
                    vertical: 1,
                    horizontal: 0,
                }),
                &mut scrollbar_state,
            );
        }
        AppMode::ExploreFiles => {
            let block = Block::bordered().title("Select File");
            let area = centered_rect(80, 60, messages_area);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);
            render_file_explorer(f, area, app);
        }
        AppMode::ShowContext => {
            let block = Block::bordered().title("Files Added to Context");
            let area = centered_rect(80, 60, messages_area);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);
            render_context_list(f, area, app);
        }
        AppMode::Notify { notification } => {
            let block = Block::bordered().title("Notification");
            let area = centered_rect(40, 40, messages_area);
            f.render_widget(Clear, area); //this clears out the background
            f.render_widget(block, area);
            render_notification(f, area, notification);
        }
    }

    let msg = match app.app_mode {
        AppMode::Editing => {
            vec![
                "Press ".into(),
                "Esc".bold(),
                " to stop editing. Press ".into(),
                "CONTROL + S (C-s)".bold(),
                " to submit the message.".into(),
            ]
        }
        AppMode::ExploreFiles => {
            vec![
                "Navigate: ".into(),
                "h/j/k/l or arrows".bold(),
                ". ".into(),
                "Enter".bold(),
                " to add file to context. ".into(),
                "d".bold(),
                " to remove from context. ".into(),
                "Esc/q".bold(),
                " to exit.".into(),
            ]
        }
        AppMode::ShowContext => {
            vec![
                "These files will be included in your next message. Press ".into(),
                "Esc/q/Enter".bold(),
                " to return.".into(),
            ]
        }
        AppMode::ModelSelection => {
            vec![
                "Navigate: ".into(),
                "j/k or Up/Down".bold(),
                ". Press ".into(),
                "Enter".bold(),
                " to select model. ".into(),
                "Esc/q".bold(),
                " to cancel.".into(),
            ]
        }
        AppMode::ShowHistory => {
            vec![
                "Navigate: ".into(),
                "j/k or Up/Down".bold(),
                ". ".into(),
                "Enter".bold(),
                " to select chat. ".into(),
                "d".bold(),
                " to delete chat. ".into(),
                "Esc/q".bold(),
                " to cancel.".into(),
            ]
        }
        AppMode::SnippetSelection => {
            vec![
                "Navigate: ".into(),
                "j/k or Up/Down".bold(),
                ". Press ".into(),
                "Enter/y".bold(),
                " to copy snippet. ".into(),
                "Esc/q".bold(),
                " to cancel.".into(),
            ]
        }
        _ => {
            vec![
                "Press ".into(),
                "Esc/q".bold(),
                " to exit. Press ".into(),
                "i".bold(),
                " to enter text. Press ".into(),
                "?".bold(),
                " for help.".into(),
            ]
        }
    };
    let text = Text::from(Line::from(msg)).patch_style(Style::default());
    let help_message = Paragraph::new(text);
    f.render_widget(help_message, help_area);

    #[cfg(not(target_os = "linux"))]
    {
        if let Some(cells) = app.selection.iter_selected_cells() {
            for (col, row) in cells {
                let cell = f.buffer_mut().cell_mut((col, row));
                // Modify the cell style to show selection
                if let Some(cell) = cell {
                    cell.set_style(SELECTED_STYLE);
                }
            }
        }

        if let Some(selected_text) = app.selection.get_selected_text(f.buffer_mut()) {
            // Trim whitespace from the selected text for each line
            let selected_text: String = selected_text
                .lines()
                .map(str::trim_end)
                .collect::<Vec<&str>>()
                .join("\n");
            app.clipboard.set_text(&selected_text).unwrap();
        }
    }
}

fn render_model_list(f: &mut Frame, area: Rect, app: &mut App) {
    let block = Block::new().padding(Padding::uniform(1));
    if app.model_list.items.is_empty() {
        let p = Paragraph::new(
            Text::from("No API keys detected, no running Ollama detected. Unable to choose model.")
                .red(),
        )
        .wrap(Wrap { trim: true })
        .block(block);
        f.render_widget(p, area);
        return;
    }
    // Iterate through all elements in the `items` and stylize them.
    let items: Vec<ListItem> = app.model_list.items.iter().map(ListItem::from).collect();

    // Create a List from all list items and highlight the currently selected one
    let list = List::new(items)
        .block(block)
        .highlight_style(SELECTED_STYLE)
        .highlight_symbol(">")
        .highlight_spacing(HighlightSpacing::Always);

    // We need to disambiguate this trait method as both `Widget` and `StatefulWidget` share the
    // same method name `render`.
    f.render_stateful_widget(list, area, &mut app.model_list.state);
}

fn render_snippet_list(f: &mut Frame, area: Rect, app: &mut App) {
    let block = Block::new().padding(Padding::uniform(1));

    // Iterate through all elements in the `items` and stylize them.
    let items: Vec<ListItem> = app
        .snippet_list
        .items
        .iter()
        .enumerate()
        .map(|(i, s)| {
            ListItem::from(format!(
                "Snippet {}: {}...",
                i + 1,
                s.text[..min(10, s.text.len())].to_owned()
            ))
        })
        .collect();

    // Create a List from all list items and highlight the currently selected one
    let list = List::new(items)
        .block(block)
        .highlight_style(SELECTED_STYLE)
        .highlight_symbol(">")
        .highlight_spacing(HighlightSpacing::Always);

    // We need to disambiguate this trait method as both `Widget` and `StatefulWidget` share the
    // same method name `render`.
    f.render_stateful_widget(list, area, &mut app.snippet_list.state);
}

fn render_chat_history_list(f: &mut Frame, area: Rect, app: &mut App) {
    let block = Block::new().padding(Padding::uniform(1));

    // Iterate through all elements in the `items` and stylize them.
    let items: Vec<ListItem> = app
        .chat_list
        .items
        .iter()
        .map(|c| ListItem::from(format!("Chat created {}", c.started_at)))
        .collect();

    // Create a List from all list items and highlight the currently selected one
    let list = List::new(items)
        .block(block)
        .highlight_style(SELECTED_STYLE)
        .highlight_symbol(">")
        .highlight_spacing(HighlightSpacing::Always);

    // We need to disambiguate this trait method as both `Widget` and `StatefulWidget` share the
    // same method name `render`.
    f.render_stateful_widget(list, area, &mut app.chat_list.state);
}

fn render_file_explorer(f: &mut Frame, area: Rect, app: &mut App) {
    let layout = Layout::horizontal([Constraint::Ratio(1, 3), Constraint::Ratio(2, 3)]);
    let file_content = get_file_content(app.file_explorer.current());

    let file_content = match file_content {
        Ok(file_content) => file_content,
        _ => "Couldn't load file.".into(),
    };

    let chunks = layout.split(area);

    f.render_widget(&app.file_explorer.widget(), chunks[0]);
    f.render_widget(Clear, chunks[1]);
    f.render_widget(
        Paragraph::new(file_content).block(
            Block::default()
                .borders(Borders::ALL)
                .border_type(BorderType::Double),
        ),
        chunks[1],
    );
}

fn render_context_list(f: &mut Frame, area: Rect, app: &mut App) {
    if let Some(context) = &app.current_context {
        let text_block = Block::new().padding(Padding::uniform(1));

        let msg: Vec<Line<'_>> = context
            .iter()
            .map(|f| Line::from(f.path().to_string_lossy()))
            .collect();
        let text = Text::from(msg).patch_style(Style::default());
        let context_text = Paragraph::new(text)
            .block(text_block)
            .wrap(Wrap { trim: true });
        f.render_widget(context_text, area);
    };
}

fn render_notification(f: &mut Frame, area: Rect, notification: &Notification) {
    let text_block = Block::new().padding(Padding::uniform(1));
    let text = match notification {
        Notification::Info(message) => Text::from(message.clone()).patch_style(Style::default()),
        Notification::Error(message) => {
            Text::from(message.clone()).patch_style(Style::default().fg(Color::Red))
        }
    };
    let context_text = Paragraph::new(text)
        .block(text_block)
        .wrap(Wrap { trim: true });
    f.render_widget(context_text, area);
}
