package crypto

import (
	"encoding/binary"
	"fmt"
	"math/big"

	"github.com/btcsuite/btcd/btcutil/base58"
)

const (
	DIGEST_LENGTH               = 5
	STATE_SIZE                  = 16
	NUM_SPLIT_AND_LOOKUP        = 4
	LOG2_STATE_SIZE             = 4
	CAPACITY                    = 6
	RATE                        = 10
	NUM_ROUNDS                  = 7
	R_MOD_P              uint64 = 4294967295
)

var (
	R, _         = new(big.Int).SetString("18446744073709551616", 10)
	LOOKUP_TABLE = [256]uint8{
		0, 7, 26, 63, 124, 215, 85, 254, 214, 228, 45, 185, 140, 173, 33, 240, 29, 177, 176, 32, 8,
		110, 87, 202, 204, 99, 150, 106, 230, 14, 235, 128, 213, 239, 212, 138, 23, 130, 208, 6, 44,
		71, 93, 116, 146, 189, 251, 81, 199, 97, 38, 28, 73, 179, 95, 84, 152, 48, 35, 119, 49, 88,
		242, 3, 148, 169, 72, 120, 62, 161, 166, 83, 175, 191, 137, 19, 100, 129, 112, 55, 221, 102,
		218, 61, 151, 237, 68, 164, 17, 147, 46, 234, 203, 216, 22, 141, 65, 57, 123, 12, 244, 54, 219,
		231, 96, 77, 180, 154, 5, 253, 133, 165, 98, 195, 205, 134, 245, 30, 9, 188, 59, 142, 186, 197,
		181, 144, 92, 31, 224, 163, 111, 74, 58, 69, 113, 196, 67, 246, 225, 10, 121, 50, 60, 157, 90,
		122, 2, 250, 101, 75, 178, 159, 24, 36, 201, 11, 243, 132, 198, 190, 114, 233, 39, 52, 21, 209,
		108, 238, 91, 187, 18, 104, 194, 37, 153, 34, 200, 143, 126, 155, 236, 118, 64, 80, 172, 89,
		94, 193, 135, 183, 86, 107, 252, 13, 167, 206, 136, 220, 207, 103, 171, 160, 76, 182, 227, 217,
		158, 56, 174, 4, 66, 109, 139, 162, 184, 211, 249, 47, 125, 232, 117, 43, 16, 42, 127, 20, 241,
		25, 149, 105, 156, 51, 53, 168, 145, 247, 223, 79, 78, 226, 15, 222, 82, 115, 70, 210, 27, 41,
		1, 170, 40, 131, 192, 229, 248, 255,
	}

	ROUND_CONSTANTS = [NUM_ROUNDS * STATE_SIZE]uint64{
		// 1st round constants
		1332676891236936200, 16607633045354064669, 12746538998793080786, 15240351333789289931,
		10333439796058208418, 986873372968378050, 153505017314310505, 703086547770691416,
		8522628845961587962, 1727254290898686320, 199492491401196126, 2969174933639985366,
		1607536590362293391, 16971515075282501568, 15401316942841283351, 14178982151025681389,
		// 2nd round constants
		2916963588744282587, 5474267501391258599, 5350367839445462659, 7436373192934779388,
		12563531800071493891, 12265318129758141428, 6524649031155262053, 1388069597090660214,
		3049665785814990091, 5225141380721656276, 10399487208361035835, 6576713996114457203,
		12913805829885867278, 10299910245954679423, 12980779960345402499, 593670858850716490,
		// 3rd round constants
		12184128243723146967, 1315341360419235257, 9107195871057030023, 4354141752578294067,
		8824457881527486794, 14811586928506712910, 7768837314956434138, 2807636171572954860,
		9487703495117094125, 13452575580428891895, 14689488045617615844, 16144091782672017853,
		15471922440568867245, 17295382518415944107, 15054306047726632486, 5708955503115886019,
		// 4th round constants
		9596017237020520842, 16520851172964236909, 8513472793890943175, 8503326067026609602,
		9402483918549940854, 8614816312698982446, 7744830563717871780, 14419404818700162041,
		8090742384565069824, 15547662568163517559, 17314710073626307254, 10008393716631058961,
		14480243402290327574, 13569194973291808551, 10573516815088946209, 15120483436559336219,
		// 5th round constants
		3515151310595301563, 1095382462248757907, 5323307938514209350, 14204542692543834582,
		12448773944668684656, 13967843398310696452, 14838288394107326806, 13718313940616442191,
		15032565440414177483, 13769903572116157488, 17074377440395071208, 16931086385239297738,
		8723550055169003617, 590842605971518043, 16642348030861036090, 10708719298241282592,
		// 6th round constants
		12766914315707517909, 11780889552403245587, 113183285481780712, 9019899125655375514,
		3300264967390964820, 12802381622653377935, 891063765000023873, 15939045541699412539,
		3240223189948727743, 4087221142360949772, 10980466041788253952, 18199914337033135244,
		7168108392363190150, 16860278046098150740, 13088202265571714855, 4712275036097525581,
		// 7th round constants
		16338034078141228133, 1455012125527134274, 5024057780895012002, 9289161311673217186,
		9401110072402537104, 11919498251456187748, 4173156070774045271, 15647643457869530627,
		15642078237964257476, 1405048341078324037, 3059193199283698832, 1605012781983592984,
		7134876918849821827, 5796994175286958720, 7251651436095127661, 4565856221886323991,
	}
	MDS_MATRIX = [STATE_SIZE][STATE_SIZE]uint64{
		{
			61402, 17845, 26798, 59689, 12021, 40901, 41351, 27521, 56951, 12034, 53865, 43244, 7454,
			33823, 28750, 1108,
		},
		{
			1108, 61402, 17845, 26798, 59689, 12021, 40901, 41351, 27521, 56951, 12034, 53865, 43244,
			7454, 33823, 28750,
		},
		{
			28750, 1108, 61402, 17845, 26798, 59689, 12021, 40901, 41351, 27521, 56951, 12034, 53865,
			43244, 7454, 33823,
		},
		{
			33823, 28750, 1108, 61402, 17845, 26798, 59689, 12021, 40901, 41351, 27521, 56951, 12034,
			53865, 43244, 7454,
		},
		{
			7454, 33823, 28750, 1108, 61402, 17845, 26798, 59689, 12021, 40901, 41351, 27521, 56951,
			12034, 53865, 43244,
		},
		{
			43244, 7454, 33823, 28750, 1108, 61402, 17845, 26798, 59689, 12021, 40901, 41351, 27521,
			56951, 12034, 53865,
		},
		{
			53865, 43244, 7454, 33823, 28750, 1108, 61402, 17845, 26798, 59689, 12021, 40901, 41351,
			27521, 56951, 12034,
		},
		{
			12034, 53865, 43244, 7454, 33823, 28750, 1108, 61402, 17845, 26798, 59689, 12021, 40901,
			41351, 27521, 56951,
		},
		{
			56951, 12034, 53865, 43244, 7454, 33823, 28750, 1108, 61402, 17845, 26798, 59689, 12021,
			40901, 41351, 27521,
		},
		{
			27521, 56951, 12034, 53865, 43244, 7454, 33823, 28750, 1108, 61402, 17845, 26798, 59689,
			12021, 40901, 41351,
		},
		{
			41351, 27521, 56951, 12034, 53865, 43244, 7454, 33823, 28750, 1108, 61402, 17845, 26798,
			59689, 12021, 40901,
		},
		{
			40901, 41351, 27521, 56951, 12034, 53865, 43244, 7454, 33823, 28750, 1108, 61402, 17845,
			26798, 59689, 12021,
		},
		{
			12021, 40901, 41351, 27521, 56951, 12034, 53865, 43244, 7454, 33823, 28750, 1108, 61402,
			17845, 26798, 59689,
		},
		{
			59689, 12021, 40901, 41351, 27521, 56951, 12034, 53865, 43244, 7454, 33823, 28750, 1108,
			61402, 17845, 26798,
		},
		{
			26798, 59689, 12021, 40901, 41351, 27521, 56951, 12034, 53865, 43244, 7454, 33823, 28750,
			1108, 61402, 17845,
		},
		{
			17845, 26798, 59689, 12021, 40901, 41351, 27521, 56951, 12034, 53865, 43244, 7454, 33823,
			28750, 1108, 61402,
		},
	}
	INIT_SPONGE     = [16]uint64{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
	INIT_SPONGE_FIX = [16]uint64{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4294967295, 4294967295, 4294967295, 4294967295, 4294967295, 4294967295}
)

func Tip5HashBelts(input []Belt) [5]uint64 {
	inputLen := len(input)
	q := inputLen / RATE
	r := inputLen % RATE

	for i := range input {
		if !BaseCheck(input[i].Value) {
			fmt.Printf("element %d must be inside the field\n", i)
		}
		input[i].Value = montify(input[i].Value)
	}

	input = append(input, Belt{Value: montify(1)})
	for i := 0; i < (RATE-r)-1; i++ {
		input = append(input, Belt{Value: montify(0)})
	}

	sponge := INIT_SPONGE
	inputAbsorb := input
	cntQ := q
	for {
		scagInput := inputAbsorb[:RATE]
		slagInput := inputAbsorb[RATE:]
		for i := 0; i < RATE; i++ {
			sponge[i] = scagInput[i].Value
		}
		sponge = permute(sponge)

		if cntQ == 0 {
			break
		}

		cntQ--
		inputAbsorb = slagInput
	}

	digest := [DIGEST_LENGTH]uint64{0, 0, 0, 0, 0}
	for i := 0; i < DIGEST_LENGTH; i++ {
		digest[i] = montReduction(new(big.Int).SetUint64(sponge[i]))
	}
	return digest
}

func Tip5HashTenCell(input []Belt) [5]uint64 {
	inputLen := len(input)
	q := inputLen / RATE
	r := inputLen % RATE
	if q != 1 || r != 0 {
		panic("input is not 10 cell")
	}

	for i := range input {
		if !BaseCheck(input[i].Value) {
			fmt.Printf("element %d must be inside the field\n", i)
		}
		input[i].Value = montify(input[i].Value)
	}

	sponge := INIT_SPONGE_FIX
	for i := 0; i < RATE; i++ {
		sponge[i] = input[i].Value
	}
	sponge = permute(sponge)
	digest := [DIGEST_LENGTH]uint64{0, 0, 0, 0, 0}
	for i := 0; i < DIGEST_LENGTH; i++ {
		digest[i] = montReduction(new(big.Int).SetUint64(sponge[i]))
	}
	return digest
}

func Tip5RehashTenCell(hash1 [5]uint64, hash2 [5]uint64) [5]uint64 {
	belts := []Belt{}
	for _, i := range hash1 {
		belts = append(belts, Belt{Value: i})
	}
	for _, i := range hash2 {
		belts = append(belts, Belt{Value: i})
	}
	return Tip5HashTenCell(belts)
}

func permute(sponge [16]uint64) [16]uint64 {
	res := sponge
	for i := 0; i < NUM_ROUNDS; i++ {
		a := sboxLayer(res)
		b := linearLayer(a)

		for j := 0; j < STATE_SIZE; j++ {
			roundConst := new(big.Int).SetUint64(ROUND_CONSTANTS[i*STATE_SIZE+j])
			product := new(big.Int).Mul(roundConst, R)
			roundConst = roundConst.Mod(product, PRIME_128)
			res[j] = badd(roundConst.Uint64(), b[j])
		}
	}
	return res
}

func Tip5HashToBase58(hash [5]uint64) string {
	result := big.NewInt(0)

	for i, value := range hash {
		result.Add(result, new(big.Int).Mul(new(big.Int).SetUint64(value), new(big.Int).Exp(PRIME_128, big.NewInt(int64(i)), nil)))
	}
	return base58.Encode(result.Bytes())
}

func Base58ToTip5Hash(data string) [5]uint64 {
	dataBigInt := new(big.Int).SetBytes(base58.Decode(data))
	result := [5]uint64{}
	for i := 0; i < 5; i++ {
		resultBigInt := new(big.Int)
		result[i] = resultBigInt.Mod(dataBigInt, PRIME_128).Uint64()
		dataBigInt.Div(dataBigInt, PRIME_128)
	}
	return result
}

func TruncGOrder(a [5]uint64) *big.Int {
	dataBigInt := new(big.Int).SetUint64(a[0])
	dataBigInt.Add(dataBigInt, new(big.Int).Mul(P_BIG, new(big.Int).SetUint64(a[1])))
	dataBigInt.Add(dataBigInt, new(big.Int).Mul(P_BIG_2, new(big.Int).SetUint64(a[2])))
	dataBigInt.Add(dataBigInt, new(big.Int).Mul(P_BIG_3, new(big.Int).SetUint64(a[3])))

	return dataBigInt.Mod(dataBigInt, G_ORDER)
}
func sboxLayer(state [16]uint64) [16]uint64 {
	res := INIT_SPONGE
	for i := 0; i < NUM_SPLIT_AND_LOOKUP; i++ {
		bytes := make([]byte, 8)
		binary.LittleEndian.PutUint64(bytes, state[i])

		for j := 0; j < 8; j++ {
			bytes[j] = LOOKUP_TABLE[bytes[j]]
		}

		res[i] = binary.LittleEndian.Uint64(bytes)
	}

	for i := NUM_SPLIT_AND_LOOKUP; i < STATE_SIZE; i++ {
		res[i] = bpow(state[i], 7)
	}
	return res
}

func linearLayer(state [16]uint64) [16]uint64 {
	res := INIT_SPONGE
	for i := 0; i < 16; i++ {
		for j := 0; j < 16; j++ {
			matrixElement := MDS_MATRIX[i][j]
			product := bmul(matrixElement, state[j])
			res[i] = badd(res[i], product)
		}
	}
	return res
}