// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
//
// SPDX-License-Identifier: GPL-2.0-or-later

package fastio

import (
	"io"
	"unicode/utf8"
)

// interfaces /////////////////////////////////////////////////////////////////

type RuneWriter interface {
	WriteRune(rune) (int, error)
}

// An AllWriter is the union of several common writer interfaces.
type AllWriter interface {
	io.Writer
	io.ByteWriter
	RuneWriter
	io.StringWriter
}

// implementations ////////////////////////////////////////////////////////////

func WriteByte(w io.Writer, b byte) error {
	var buf [1]byte
	buf[0] = b
	_, err := w.Write(buf[:])
	return err
}

func WriteRune(w io.Writer, r rune) (int, error) {
	var buf [utf8.UTFMax]byte
	n := utf8.EncodeRune(buf[:], r)
	return w.Write(buf[:n])
}

func WriteString(w io.Writer, s string) (int, error) {
	return w.Write([]byte(s))
}

// wrappers ///////////////////////////////////////////////////////////////////

// NNN

type (
	writerNNN        interface{ io.Writer }
	writerNNNWrapper struct{ writerNNN }
)

func (w writerNNNWrapper) WriteByte(b byte) error            { return WriteByte(w, b) }
func (w writerNNNWrapper) WriteRune(r rune) (int, error)     { return WriteRune(w, r) }
func (w writerNNNWrapper) WriteString(s string) (int, error) { return WriteString(w, s) }

// NNY

type (
	writerNNY interface {
		io.Writer
		io.StringWriter
	}
	writerNNYWrapper struct{ writerNNY }
)

func (w writerNNYWrapper) WriteByte(b byte) error        { return WriteByte(w, b) }
func (w writerNNYWrapper) WriteRune(r rune) (int, error) { return WriteRune(w, r) }

// NYN

type (
	writerNYN interface {
		io.Writer
		RuneWriter
	}
	writerNYNWrapper struct{ writerNYN }
)

func (w writerNYNWrapper) WriteByte(b byte) error            { return WriteByte(w, b) }
func (w writerNYNWrapper) WriteString(s string) (int, error) { return WriteString(w, s) }

// NYY

type (
	writerNYY interface {
		io.Writer
		RuneWriter
		io.StringWriter
	}
	writerNYYWrapper struct{ writerNYY }
)

func (w writerNYYWrapper) WriteByte(b byte) error { return WriteByte(w, b) }

// YNN

type (
	writerYNN interface {
		io.Writer
		io.ByteWriter
	}
	writerYNNWrapper struct{ writerYNN }
)

func (w writerYNNWrapper) WriteRune(r rune) (int, error)     { return WriteRune(w, r) }
func (w writerYNNWrapper) WriteString(s string) (int, error) { return WriteString(w, s) }

// YNY

type (
	writerYNY interface {
		io.Writer
		io.ByteWriter
		io.StringWriter
	}
	writerYNYWrapper struct{ writerYNY }
)

func (w writerYNYWrapper) WriteRune(r rune) (int, error) { return WriteRune(w, r) }

// YYN

type (
	writerYYN interface {
		io.Writer
		io.ByteWriter
		RuneWriter
	}
	writerYYNWrapper struct{ writerYYN }
)

func (w writerYYNWrapper) WriteString(s string) (int, error) { return WriteString(w, s) }

// NewAllWriter wraps an io.Writer turning it in to an AllWriter.  If
// the io.Writer already has any of the other write methods, then its
// native version of those methods are used.
func NewAllWriter(inner io.Writer) AllWriter {
	switch inner := inner.(type) {
	// 3 Y bits
	case AllWriter: // YYY:
		return inner
	// 2 Y bits
	case writerNYY:
		return writerNYYWrapper{writerNYY: inner}
	case writerYNY:
		return writerYNYWrapper{writerYNY: inner}
	case writerYYN:
		return writerYYNWrapper{writerYYN: inner}
	// 1 Y bit
	case writerNNY:
		return writerNNYWrapper{writerNNY: inner}
	case writerNYN:
		return writerNYNWrapper{writerNYN: inner}
	case writerYNN:
		return writerYNNWrapper{writerYNN: inner}
	// 0 Y bits
	default: // NNN:
		return writerNNNWrapper{writerNNN: inner}
	}
}

// discard /////////////////////////////////////////////////////////////////////

// Discard is like io.Discard, but implements AllWriter.
var Discard = discard{}

type discard struct{}

func (discard) Write(p []byte) (int, error)       { return len(p), nil }
func (discard) WriteByte(b byte) error            { return nil }
func (discard) WriteRune(r rune) (int, error)     { return 0, nil }
func (discard) WriteString(s string) (int, error) { return len(s), nil }