package cryptosquare
import (
	"strings"
	"testing"
)
var tests = []struct {
	pt string // plain text
	ct string // cipher text
}{
	{
		"s#$%^&plunk",
		"su pn lk",
	},
	{
		"1, 2, 3 GO!",
		"1g 2o 3 ",
	},
	{
		"1234",
		"13 24",
	},
	{
		"123456789",
		"147 258 369",
	},
	{
		"123456789abc",
		"159 26a 37b 48c",
	},
	{
		"Never vex thine heart with idle woes",
		"neewl exhie vtetw ehaho ririe vntds",
	},
	{
		"ZOMG! ZOMBIES!!!",
		"zzi ooe mms gb ",
	},
	{
		"Time is an illusion. Lunchtime doubly so.",
		"tasney inicds miohoo elntu  illib  suuml ",
	},
	{
		"We all know interspecies romance is weird.",
		"wneiaw eorene awssci liprer lneoid ktcms ",
	},
	{
		"Madness, and then illumination.",
		"msemo aanin dnin  ndla  etlt  shui ",
	},
	{
		"Vampires are people too!",
		"vrel aepe mset paoo irpo",
	},
	{
		"",
		"",
	},
	{
		"1",
		"1",
	},
	{
		"12",
		"1 2",
	},
	{
		"12 3",
		"13 2 ",
	},
	{
		"12345678",
		"147 258 36 ",
	},
	{
		"123456789a",
		"159 26a 37  48 ",
	},
	{
		"If man was meant to stay on the ground god would have given us roots",
		"imtgdvs fearwer mayoogo anouuio ntnnlvt wttddes aohghn  sseoau ",
	},
	{
		"Have a nice day. Feed the dog & chill out!",
		"hifei acedl veeol eddgo aatcu nyhht",
	},
}
func TestEncode(t *testing.T) {
	for _, test := range tests {
		if ct := Encode(test.pt); ct != test.ct {
			t.Fatalf(`Encode(%q):
got  %q
want %q`, test.pt, ct, test.ct)
		}
	}
}
func TestDecode(t *testing.T) {
	for _, tc := range tests {
		want := strings.Map(normalize, tc.pt)
		want = strings.ToLower(want)
		if dt := Decode(tc.ct); dt != want {
			t.Fatalf(`Decode(%q):
got  %q
want %q`, tc.ct, dt, want)
		}
	}
}
func BenchmarkEncode(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range tests {
			Encode(test.pt)
		}
	}
}

// Package cryptosquare is for handling crypto squares.
package cryptosquare
import (
	"math"
	"strings"
	"unicode"
)
// Encode encodes in with cryptosquare.
func Encode(in string) string {
	in = strings.Map(normalize, in)
	in = strings.ToLower(in)
	// calculate row and column number
	r := int(math.Sqrt(float64(len(in))))
	c := r
	for ; c*r < len(in); r, c = r+1, c+1 {
		if (c+1)*r >= len(in) {
			c++
			break
		}
	}
	// encode the input
	out := strings.Builder{}
	for j := 0; j < c; j++ {
		for i := j; i < len(in); i += c {
			out.WriteByte(in[i])
		}
		if (len(out.String())-j)%r != 0 {
			out.WriteRune(' ')
		}
		if j != c-1 {
			out.WriteRune(' ')
		}
	}
	return out.String()
}
func normalize(c rune) rune {
	if unicode.IsPunct(c) ||
		unicode.IsSymbol(c) ||
		unicode.IsSpace(c) {
		return -1 // negative return value removes the rune
	}
	return c
}
// Decode decodes cryptosquare encoded string. Because encoding is lossy, it will return
// the string without spaces, punctuation and symbols in lowercase.
func Decode(in string) string {
	c := strings.IndexRune(in, ' ') + 1
	if c == 0 {
		return in
	}
	out := strings.Builder{}
	for j := 0; j < c; j++ {
		for i := j; i < len(in); i += c {
			if in[i] != ' ' {
				out.WriteByte(in[i])
			}
		}
	}
	return out.String()
}