1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

/// Write colored text to the screen
#[derive(Debug)]
pub struct Console<S>
where
    S: crate::WinconStream + std::io::Write,
{
    stream: Option<S>,
    initial_fg: Option<anstyle::AnsiColor>,
    initial_bg: Option<anstyle::AnsiColor>,
    last_fg: Option<anstyle::AnsiColor>,
    last_bg: Option<anstyle::AnsiColor>,
}

impl<S> Console<S>
where
    S: crate::WinconStream + std::io::Write,
{
    pub fn new(stream: S) -> Result<Self, S> {
        // HACK: Assuming the error from `get_colors()` will be present on `write` and doing basic
        // ops on the stream will cause the same result
        let (initial_fg, initial_bg) = match stream.get_colors() {
            Ok(ok) => ok,
            Err(_) => {
                return Err(stream);
            }
        };
        Ok(Self {
            stream: Some(stream),
            initial_fg,
            initial_bg,
            last_fg: initial_fg,
            last_bg: initial_bg,
        })
    }

    /// Write colored text to the screen
    pub fn write(
        &mut self,
        fg: Option<anstyle::AnsiColor>,
        bg: Option<anstyle::AnsiColor>,
        data: &[u8],
    ) -> std::io::Result<usize> {
        self.apply(fg, bg)?;
        let written = self.as_stream_mut().write(data)?;
        Ok(written)
    }

    pub fn flush(&mut self) -> std::io::Result<()> {
        self.as_stream_mut().flush()
    }

    /// Change the terminal back to the initial colors
    pub fn reset(&mut self) -> std::io::Result<()> {
        self.apply(self.initial_fg, self.initial_bg)
    }

    /// Close the stream, reporting any errors
    pub fn close(mut self) -> std::io::Result<()> {
        self.reset()
    }

    /// Allow changing the stream
    pub fn map<S1: crate::WinconStream + std::io::Write>(
        mut self,
        op: impl FnOnce(S) -> S1,
    ) -> Console<S1> {
        Console {
            stream: Some(op(self.stream.take().unwrap())),
            initial_fg: self.initial_fg,
            initial_bg: self.initial_bg,
            last_fg: self.last_fg,
            last_bg: self.last_bg,
        }
    }

    /// Get the inner writer
    #[inline]
    pub fn into_inner(mut self) -> S {
        let _ = self.reset();
        self.stream.take().unwrap()
    }

    fn apply(
        &mut self,
        fg: Option<anstyle::AnsiColor>,
        bg: Option<anstyle::AnsiColor>,
    ) -> std::io::Result<()> {
        let fg = fg.or(self.initial_fg);
        let bg = bg.or(self.initial_bg);
        if fg == self.last_fg && bg == self.last_bg {
            return Ok(());
        }

        // Ensure everything is written with the last set of colors before applying the next set
        self.as_stream_mut().flush()?;

        self.as_stream_mut().set_colors(fg, bg)?;
        self.last_fg = fg;
        self.last_bg = bg;

        Ok(())
    }

    fn as_stream_mut(&mut self) -> &mut S {
        self.stream.as_mut().unwrap()
    }
}

impl<S> Console<S>
where
    S: crate::WinconStream + std::io::Write,
    S: crate::Lockable,
    <S as crate::Lockable>::Locked: crate::WinconStream + std::io::Write,
{
    /// Get exclusive access to the `Console`
    ///
    /// Why?
    /// - Faster performance when writing in a loop
    /// - Avoid other threads interleaving output with the current thread
    #[inline]
    pub fn lock(mut self) -> <Self as crate::Lockable>::Locked {
        Console {
            stream: Some(self.stream.take().unwrap().lock()),
            initial_fg: self.initial_fg,
            initial_bg: self.initial_bg,
            last_fg: self.last_fg,
            last_bg: self.last_bg,
        }
    }
}

impl<S> crate::Lockable for Console<S>
where
    S: crate::WinconStream + std::io::Write,
    S: crate::Lockable,
    <S as crate::Lockable>::Locked: crate::WinconStream + std::io::Write,
{
    type Locked = Console<<S as crate::Lockable>::Locked>;

    #[inline]
    fn lock(self) -> Self::Locked {
        self.lock()
    }
}

impl<S> Drop for Console<S>
where
    S: crate::WinconStream + std::io::Write,
{
    fn drop(&mut self) {
        // Otherwise `Console::lock` took it
        if self.stream.is_some() {
            let _ = self.reset();
        }
    }
}