use std::io::Write;
use textwrap::fill as textwrap_fill;

fn termwidth() -> usize {
	use terminal_size::{terminal_size, Height, Width};
	let size = terminal_size();
	if let Some((Width(w), Height(_))) = size {
		std::cmp::min(w as usize, 80)
	} else {
		80
	}
}

fn fill(str: &str) -> Option<String> {
	let width = termwidth();
	let filled = textwrap_fill(str, width);
	if filled.contains('\n') {
		Some(filled)
	} else {
		None
	}
}

fn clear_format(str: &str) -> String {
	let width = termwidth();
	let whitespace = " ".repeat(width);
	let filled = textwrap_fill(str, width);
	format!("\r{}\r{}", whitespace, filled)
}

use colored::Colorize;

#[derive(PartialEq)]
enum State {
	Write,
	Think,
	Buf,
}

pub struct Buffer {
	pub buf: String,
	state: State,
}

impl Buffer {
	pub fn new() -> Self {
		Buffer {
			buf: String::new(),
			state: State::Write,
		}
	}
	pub fn proc(&mut self, data: &str) {
		match self.state {
			State::Write => self.proc_write(data),
			State::Think => self.proc_think(data),
			State::Buf => self.buf.push_str(data),
		}
	}

	pub fn print_return_remain(&mut self) -> String {
		let buffered = self.buf.trim().to_string();
		self.buf.clear();
		if self.state == State::Buf {
			return buffered;
		}

		let split = buffered.split_once("<suggestions>");
		if let Some((first, last)) = split {
			eprint!("{}", first);
			std::io::stdout().flush().unwrap();
			return last.to_string();
		}
		"".to_string()
	}

	fn proc_write(&mut self, data: &str) {
		if !data.contains("\n") {
			self.buf.push_str(data);
			let buffered = self.buf.trim().to_string();
			let filled = fill(&buffered);
			if let Some(filled) = filled {
				self.buf.clear();
				let formatted = clear_format(&filled);
				eprint!("{}", formatted);
				self.buf.push_str(formatted.split_once("\n").unwrap().1);
				std::io::stdout().flush().unwrap();
				return;
			}
			eprint!("{}", data);
			std::io::stdout().flush().unwrap();
			return;
		}

		let mut data = data.to_string();
		while data.contains("\n") {
			let lines = data.split_once("\n").unwrap();
			let first = lines.0;
			let last = lines.1;
			self.buf.push_str(first);
			let buffered = self.buf.trim().to_string();
			self.buf.clear();
			if buffered.ends_with("<note>") {
				let warn = format!("{}:", t!("ai-suggestion"))
					.bold()
					.blue()
					.to_string();
				let first = buffered.replace("<note>", &warn);
				eprintln!("\r{}", first);
				std::io::stdout().flush().unwrap();
			} else if buffered.ends_with("</note>") {
				let tag = "</note>";
				let whitespace = " ".repeat(tag.len());
				let clear = whitespace.to_string();
				let first = buffered.replace("</note>", &clear);
				eprintln!("\r{}", first);
				self.state = State::Buf;
				std::io::stdout().flush().unwrap();
			} else if buffered.ends_with("<think>") {
				let tag = "<think>";
				let warn = format!("{}:", t!("ai-thinking")).bold().blue().to_string();
				let first = buffered.replace(tag, &warn);
				self.state = State::Think;
				eprintln!("\r{}", first);
				std::io::stdout().flush().unwrap();
			} else if buffered.ends_with("</think>") {
				let tag = "</think>";
				let whitespace = " ".repeat(tag.len());
				let clear = whitespace.to_string();
				let first = buffered.replace(tag, &clear);
				eprintln!("\r{}", first);
				std::io::stdout().flush().unwrap();
			} else if buffered.ends_with("```") {
				let tag = "```";
				let whitespace = " ".repeat(tag.len());
				let formatted = format!("\r{}", whitespace);
				let first = buffered.replace(tag, &formatted);
				eprintln!("{}", first);
				std::io::stdout().flush().unwrap();
			} else {
				eprintln!("{}", first);
				std::io::stdout().flush().unwrap();
			}
			data = last.to_string();
		}
		eprint!("{}", data);
	}

	fn proc_think(&mut self, data: &str) {
		if !data.contains("\n") {
			self.buf.push_str(data);
			let buffered = self.buf.trim().to_string();
			let filled = fill(&buffered);
			if let Some(filled) = filled {
				self.buf.clear();
				let formatted = clear_format(&filled);
				eprint!("{}", formatted);
				self.buf.push_str(formatted.split_once("\n").unwrap().1);
				std::io::stdout().flush().unwrap();
				return;
			}
			eprint!("{}", data);
			std::io::stdout().flush().unwrap();
			return;
		}

		let mut data = data.to_string();
		while data.contains("\n") {
			let lines = data.split_once("\n").unwrap();
			let first = lines.0;
			let last = lines.1;
			self.buf.push_str(first);
			let buffered = self.buf.trim().to_string();
			self.buf.clear();
			if buffered.ends_with("</think>") {
				let tag = "</think>";
				let whitespace = " ".repeat(tag.len());
				let clear = whitespace.to_string();
				let first = buffered.replace(tag, &clear);
				self.state = State::Write;
				eprintln!("\r{}", first);
				std::io::stdout().flush().unwrap();
			} else {
				eprintln!("{}", first);
				std::io::stdout().flush().unwrap();
			}
			data = last.to_string();
		}
		eprint!("{}", data);
	}
}
