use crate::markdown::markdown;
use crate::text::{clean, wrap, ACCENT, ERROR, MUTED};
use klbr_ipc::v2::{CellState, CellTrace};
use ratatui::{
    style::{Color, Modifier, Style},
    text::{Line, Span},
};
use std::collections::HashMap;

#[derive(Clone, Debug)]
pub enum Entry {
    Text {
        label: String,
        body: String,
        muted: bool,
    },
    Cell {
        code: String,
        state: Option<CellState>,
        output: String,
        error: Option<String>,
        truncated: usize,
        trace: Option<CellTrace>,
    },
    Note {
        text: String,
        error: bool,
    },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Format {
    pub width: usize,
    pub expanded: bool,
    pub thinking: bool,
}

fn label(text: impl Into<String>, color: Color) -> Line<'static> {
    Line::styled(
        text.into(),
        Style::default().fg(color).add_modifier(Modifier::BOLD),
    )
}

fn preview(
    mut rows: Vec<Line<'static>>,
    limit: usize,
    expanded: bool,
    tail: bool,
) -> Vec<Line<'static>> {
    if expanded || rows.len() <= limit {
        return rows;
    }
    let hidden = rows.len() - limit;
    let hint = Line::styled(
        format!(
            "… {hidden} {} hidden · /v expand",
            if hidden == 1 { "row" } else { "rows" }
        ),
        Style::default().fg(MUTED),
    );
    if tail {
        rows.drain(..hidden);
        rows.insert(0, hint);
    } else {
        rows.truncate(limit);
        rows.push(hint);
    }
    rows
}

impl Entry {
    pub fn rows(&self, format: Format) -> Vec<Line<'static>> {
        let mut rows = match self {
            Self::Text {
                label: name,
                body,
                muted,
            } => {
                let style = if name == "you" {
                    Style::default().fg(ACCENT)
                } else if *muted {
                    Style::default().fg(MUTED)
                } else {
                    Style::default().add_modifier(Modifier::BOLD)
                };
                let mut rows = vec![Line::styled(clean(name), style)];
                rows.extend(
                    if matches!(name.as_str(), "thinking" | "scratchpad") && !format.thinking {
                        vec![Line::styled("hidden · /t show", Style::default().fg(MUTED))]
                    } else {
                        wrap(markdown(body, *muted, format.width), format.width)
                    },
                );
                rows
            }
            Self::Note { text, error } => wrap(
                vec![Line::styled(
                    clean(text),
                    Style::default().fg(if *error { ERROR } else { MUTED }),
                )],
                format.width,
            ),
            // A start is retained for correlation/backfill, but only the transient action is drawn.
            Self::Cell { state: None, .. } => return Vec::new(),
            Self::Cell {
                code,
                state,
                output,
                error,
                truncated,
                trace,
            } => {
                let failed = matches!(state, Some(CellState::Error | CellState::Interrupted));
                let has_trace = trace
                    .as_ref()
                    .is_some_and(|t| !t.activities.is_empty() || !t.changes.is_empty());
                let mut rows = Vec::new();
                if !has_trace || format.expanded || failed {
                    let count = code.split('\n').count();
                    let suffix = match state {
                        Some(CellState::Error) => " · failed",
                        Some(CellState::Interrupted) => " · interrupted",
                        Some(CellState::Yielded) => " · yielded",
                        _ => "",
                    };
                    rows.push(Line::styled(
                        format!(
                            "python · {count} {}{suffix}",
                            if count == 1 { "line" } else { "lines" }
                        ),
                        Style::default().fg(if failed { ERROR } else { MUTED }),
                    ));
                }
                if let Some(trace) = trace {
                    if !trace.activities.is_empty() {
                        rows.push(label(
                            if trace.activities.iter().any(|a| a.kind == "run") {
                                "executed"
                            } else {
                                "explored"
                            },
                            Color::Reset,
                        ));
                        for activity in &trace.activities {
                            rows.extend(wrap(
                                vec![Line::from(vec![
                                    Span::raw("  "),
                                    Span::styled(
                                        clean(&activity.kind),
                                        Style::default().fg(Color::Cyan),
                                    ),
                                    Span::raw(format!(" {}", clean(&activity.target))),
                                ])],
                                format.width,
                            ));
                        }
                    }
                    for change in &trace.changes {
                        let unavailable = change.truncated && change.diff.is_empty();
                        let mut heading = vec![Span::styled(
                            format!("edited {}", clean(&change.path)),
                            Style::default().add_modifier(Modifier::BOLD),
                        )];
                        if !unavailable {
                            heading.extend([
                                Span::styled(
                                    format!("  +{}", change.added),
                                    Style::default().fg(Color::Green),
                                ),
                                Span::styled(
                                    format!(" −{}", change.removed),
                                    Style::default().fg(Color::Red),
                                ),
                            ]);
                        }
                        rows.push(Line::from(heading));
                        let diff = clean(&change.diff)
                            .split('\n')
                            .map(|line| {
                                Line::styled(
                                    line.to_owned(),
                                    Style::default().fg(if line.starts_with('+') {
                                        Color::Green
                                    } else if line.starts_with('-') {
                                        Color::Red
                                    } else {
                                        MUTED
                                    }),
                                )
                            })
                            .collect();
                        rows.extend(preview(wrap(diff, format.width), 8, format.expanded, false));
                        if change.truncated {
                            rows.push(Line::styled(
                                "diff truncated at source",
                                Style::default().fg(MUTED),
                            ));
                        }
                    }
                    if trace.truncated {
                        rows.push(Line::styled(
                            "activity capture limited; some operations are not shown",
                            Style::default().fg(MUTED),
                        ));
                    }
                }
                if format.expanded {
                    rows.extend(wrap(
                        markdown(&format!("```python\n{code}\n```"), false, format.width),
                        format.width,
                    ));
                }
                if format.expanded || failed || !has_trace {
                    let mut result = output.trim_end().to_owned();
                    if let Some(error) = error {
                        if !result.is_empty() {
                            result.push('\n');
                        }
                        result.push_str(error.trim_end());
                    }
                    if result.is_empty() {
                        result = "(no output)".into();
                    }
                    let output = wrap(
                        clean(&result)
                            .split('\n')
                            .map(|s| {
                                Line::styled(
                                    s.to_owned(),
                                    if failed {
                                        Style::default().fg(ERROR)
                                    } else {
                                        Style::default()
                                    },
                                )
                            })
                            .collect(),
                        format.width,
                    );
                    rows.extend(preview(output, 3, format.expanded, failed));
                }
                if *truncated > 0 {
                    rows.push(Line::styled(
                        format!("{truncated} output bytes omitted by runtime"),
                        Style::default().fg(MUTED),
                    ));
                }
                rows
            }
        };
        rows = wrap(rows, format.width);
        rows.push(Line::default());
        rows
    }
}

struct Block {
    key: String,
    seq: i64,
    revision: u64,
    entry: Entry,
}
#[derive(Default)]
struct Layout {
    format: Option<Format>,
    rows: Vec<Vec<Line<'static>>>,
    revisions: Vec<u64>,
    ends: Vec<usize>,
    dirty_from: usize,
}

#[derive(Default)]
pub struct Transcript {
    blocks: Vec<Block>,
    lookup: HashMap<String, usize>,
    layouts: Vec<Layout>,
    sorted: bool,
}

impl Transcript {
    pub fn upsert(&mut self, key: String, seq: i64, entry: Entry) {
        if let Some(&index) = self.lookup.get(&key) {
            for layout in &mut self.layouts {
                layout.dirty_from = layout.dirty_from.min(index);
            }
            let block = &mut self.blocks[index];
            if seq < block.seq {
                block.seq = seq;
                self.sorted = false;
            }
            block.entry = entry;
            block.revision += 1;
        } else {
            for layout in &mut self.layouts {
                layout.dirty_from = layout.dirty_from.min(self.blocks.len());
            }
            if self.blocks.is_empty() {
                self.sorted = true;
            }
            if self.blocks.last().is_some_and(|last| last.seq > seq) {
                self.sorted = false;
            }
            self.lookup.insert(key.clone(), self.blocks.len());
            self.blocks.push(Block {
                key,
                seq,
                entry,
                revision: 1,
            });
        }
    }

    pub fn entry_mut(&mut self, key: &str) -> Option<&mut Entry> {
        let index = *self.lookup.get(key)?;
        for layout in &mut self.layouts {
            layout.dirty_from = layout.dirty_from.min(index);
        }
        let block = &mut self.blocks[index];
        block.revision += 1;
        Some(&mut block.entry)
    }

    pub fn entry(&self, key: &str) -> Option<&Entry> {
        Some(&self.blocks[*self.lookup.get(key)?].entry)
    }

    pub fn remove(&mut self, key: &str) {
        if let Some(index) = self.lookup.remove(key) {
            self.blocks.remove(index);
            self.sorted = false;
            self.lookup = self
                .blocks
                .iter()
                .enumerate()
                .map(|(i, block)| (block.key.clone(), i))
                .collect();
        }
    }

    fn layout(&mut self, format: Format) -> &Layout {
        if !self.sorted {
            self.blocks.sort_by_key(|block| block.seq);
            self.lookup = self
                .blocks
                .iter()
                .enumerate()
                .map(|(i, block)| (block.key.clone(), i))
                .collect();
            self.layouts.clear();
            self.sorted = true;
        }
        let index = self
            .layouts
            .iter()
            .position(|layout| layout.format == Some(format));
        let mut layout = index
            .map(|i| self.layouts.remove(i))
            .unwrap_or_else(|| Layout {
                format: Some(format),
                ..Layout::default()
            });
        // Only changed blocks are formatted. Offscreen rows are never cloned into a frame.
        let mut changed = false;
        for (i, block) in self.blocks.iter().enumerate().skip(layout.dirty_from) {
            if layout.revisions.get(i) != Some(&block.revision) {
                let rows = block.entry.rows(format);
                if i == layout.rows.len() {
                    layout.rows.push(rows);
                    layout.revisions.push(block.revision);
                    layout.ends.push(0);
                } else {
                    layout.rows[i] = rows;
                    layout.revisions[i] = block.revision;
                }
                changed = true;
            }
            if changed {
                layout.ends[i] =
                    layout.ends.get(i.wrapping_sub(1)).copied().unwrap_or(0) + layout.rows[i].len();
            }
        }
        layout.dirty_from = usize::MAX;
        self.layouts.insert(0, layout);
        self.layouts.truncate(2);
        &self.layouts[0]
    }

    pub fn len(&mut self, format: Format) -> usize {
        self.layout(format).ends.last().copied().unwrap_or(0)
    }

    pub fn window(&mut self, format: Format, top: usize, height: usize) -> Vec<Line<'static>> {
        let layout = self.layout(format);
        let first = layout.ends.partition_point(|end| *end <= top);
        let mut offset =
            top.saturating_sub(layout.ends.get(first.wrapping_sub(1)).copied().unwrap_or(0));
        let mut rows = Vec::with_capacity(height);
        for block in &layout.rows[first..] {
            for row in block.iter().skip(offset).take(height - rows.len()) {
                rows.push(row.clone());
            }
            if rows.len() == height {
                break;
            }
            offset = 0;
        }
        rows
    }

    pub fn anchor(&mut self, format: Format, row: usize) -> Option<(String, usize)> {
        let layout = self.layout(format);
        let index = layout.ends.partition_point(|end| *end <= row);
        let offset =
            row.saturating_sub(layout.ends.get(index.wrapping_sub(1)).copied().unwrap_or(0));
        self.blocks
            .get(index)
            .map(|block| (block.key.clone(), offset))
    }

    pub fn locate(&mut self, format: Format, anchor: &(String, usize)) -> Option<usize> {
        self.layout(format);
        let index = *self.lookup.get(&anchor.0)?;
        let layout = &self.layouts[0];
        Some(
            layout.ends.get(index.wrapping_sub(1)).copied().unwrap_or(0)
                + anchor.1.min(layout.rows[index].len().saturating_sub(1)),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    fn format(width: usize) -> Format {
        Format {
            width,
            expanded: false,
            thinking: true,
        }
    }
    #[test]
    fn previews_bound_display_rows_and_keep_errors() {
        let entry = Entry::Cell {
            code: "print(1)".into(),
            state: Some(CellState::Error),
            output: "x".repeat(10000),
            error: Some("SyntaxError: broken".into()),
            truncated: 0,
            trace: None,
        };
        let rows = entry.rows(format(40));
        assert!(rows.len() < 12);
        assert!(rows
            .iter()
            .any(|row| row.to_string().contains("SyntaxError")));
        assert!(rows
            .iter()
            .any(|row| row.to_string().contains("1 line · failed")));
    }
    #[test]
    fn history_and_reflow_keep_the_same_record_anchor() {
        let mut transcript = Transcript::default();
        for i in 10..110 {
            transcript.upsert(
                i.to_string(),
                i,
                Entry::Text {
                    label: "you".into(),
                    body: format!("row {i} {}", "context ".repeat(30)),
                    muted: false,
                },
            );
        }
        let anchor = transcript.anchor(format(80), 50).unwrap();
        for i in 0..10 {
            transcript.upsert(
                i.to_string(),
                i,
                Entry::Note {
                    text: format!("older {i}"),
                    error: false,
                },
            );
        }
        let position = transcript.locate(format(40), &anchor).unwrap();
        assert_eq!(transcript.anchor(format(40), position).unwrap().0, anchor.0);
        let rows = transcript.window(format(40), position, 15);
        assert_eq!(rows.len(), 15);
        assert!(rows.iter().all(|row| row.width() <= 40));
    }
}