package wrapwriter

import (
	"bytes"
	"io"
	"math"
)

var (
	nl = []byte{'\n'}
	sp = []byte{' '}
)

type Writer struct {
	width int
	buf   []byte
	bol   bool
	w     io.Writer
}

func NewWriter(w io.Writer, width int) *Writer {
	return &Writer{width: width, w: w, bol: true}
}

func (w *Writer) Write(p []byte) (n int, err error) {
	for _, c := range p {
		w.buf = append(w.buf, c)
		if c == '\n' && w.bol {
			if err = w.wrap(); err != nil {
				return 0, err
			}
			if _, err = w.w.Write(nl); err != nil {
				return 0, err
			}
		}
		w.bol = c == '\n'
	}
	return len(p), nil
}

func (w *Writer) Flush() error {
	return w.wrap()
}

func (w *Writer) wrap() (err error) {
	b := bytes.TrimSpace(w.buf)
	w.buf = nil
	if len(b) < 1 {
		return nil
	}
	words := bytes.Split(bytes.Replace(b, nl, sp, -1), sp)
	for _, line := range wrapWords(words, w.width) {
		for i, word := range line {
			if _, err = w.w.Write(word); err != nil {
				return err
			}
			if i == len(line)-1 {
				_, err = w.w.Write(nl)
			} else {
				_, err = w.w.Write(sp)
			}
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// wrapWords is the low-level line-breaking algorithm, useful if you need more
// control over the details of the text wrapping process. For most uses, either
// Wrap or WrapBytes will be sufficient and more convenient. 
//
// wrapWords splits a list of words into lines with minimal "raggedness",
// treating each byte as one unit, accounting for 1 unit between adjacent
// words on each line, and attempting to limit lines to lim units. Raggedness
// is the total error over all lines, where error is the square of the
// difference of the length of the line and lim. Too-long lines (which only
// happen when a single word is longer than lim units) have lim**2 penalty
// units added to the error.
// TODO(kr): count runes instead of bytes for width
func wrapWords(words [][]byte, lim int) [][][]byte {
	pen := lim * lim
	n := len(words)
	length := make([][]int, n)
	for i := 0; i < n; i++ {
		length[i] = make([]int, n)
		length[i][i] = len(words[i])
		for j := i + 1; j < n; j++ {
			length[i][j] = length[i][j-1] + 1 + len(words[j])
		}
	}

	nbrk := make([]int, n)
	cost := make([]int, n)
	for i := range cost {
		cost[i] = math.MaxInt32
	}
	for i := n - 1; i >= 0; i-- {
		if length[i][n-1] <= lim {
			cost[i] = 0
			nbrk[i] = n
		} else {
			for j := i + 1; j < n; j++ {
				d := lim - length[i][j-1]
				c := d*d + cost[j]
				if length[i][j-1] > lim {
					c += pen // too-long lines get a worse penalty
				}
				if c < cost[i] {
					cost[i] = c
					nbrk[i] = j
				}
			}
		}
	}

	var lines [][][]byte
	i := 0
	for i < n {
		lines = append(lines, words[i:nbrk[i]])
		i = nbrk[i]
	}
	return lines
}