package decrypt

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/cloudflare/circl/kem/kyber/kyber768"
	"golang.org/x/crypto/argon2"
	cc20 "golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/hkdf"
	"sources.truenas.cloud/code/syscrypt"
	"sources.truenas.cloud/code/syscrypt/internal/utils"
	"sources.truenas.cloud/code/syscrypt/internal/vars"
)

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func Validate() {

	utils.ClearTerminal()

	decryptAllowedFlags := map[string]struct{}{
		"-C": {},
		"-k": {},
		"-i": {},
		"-c": {},
		"-o": {},
		"-A": {},
		"-P": {},
	}

	utils.ValidateAllowedFlags(decryptAllowedFlags)

	decryptRequiredFlags := map[string]bool{
		"-k": true,
		"-i": true,
	}

	_ = decryptRequiredFlags

	//utils.ValidateRequiredFlags(decryptRequiredFlags, "decrypt")

	// -- Keys

	isKeySet, keyHasValue := utils.IsFlagPassed("k")
	keyValue, _ := utils.GetFlagValue("k")
	keyIsValidPath := utils.IsValidPath(keyValue)
	keyExists := utils.FileExists(keyValue)
	keyIsValidJson, _ := utils.ValidateJSON(keyValue)

	if isKeySet && !keyHasValue {
		msg := fmt.Sprintf("%s: -k KEY: requires a value.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if isKeySet && !keyIsValidPath {
		msg := fmt.Sprintf("%s: -k KEY: requires a valid file path.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if isKeySet && !keyExists {
		msg := fmt.Sprintf("%s: -k KEY: key file does not exist.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if !keyIsValidJson {
		msg := fmt.Sprintf("%s: -k KEY: Invalid JSON format.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	var privateKeyPrefix string

	privateKey, _ := utils.FetchFileKey(keyValue, "key")
	privateKeyParts := strings.Split(privateKey, "--")
	if len(privateKeyParts) != 2 {
		msg := fmt.Sprintf("%s: -k KEY: Mismatch: Invalid Prefix '%s'. Invalid Length.\n", vars.CommandFlag, privateKeyParts[0]+vars.EndAnchor)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	privateKeyStart := privateKeyParts[0]
	privateKeyPrefix = strings.ReplaceAll(vars.PrivateKeyPrefixLabel, "--", "")

	var match bool
	if privateKeyStart == privateKeyPrefix {
		match = true
	} else {
		match = false
	}

	if !match {
		msg := fmt.Sprintf("%s: -k KEY: Mismatch: Invalid Prefix '%s'\n", vars.CommandFlag, privateKeyStart+vars.EndAnchor)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	// -- Input

	isInputSet, inputHasValue := utils.IsFlagPassed("i")
	inputValue, _ := utils.GetFlagValue("i")
	inputIsValidPath := utils.IsValidPath(inputValue)
	inputExists := utils.FileExists(inputValue)

	if isInputSet && !inputHasValue {
		msg := fmt.Sprintf("%s: -i INPUT: Requires a value.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if isInputSet && !inputIsValidPath {
		msg := fmt.Sprintf("%s: -i INPUT: Requires a valid file path.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if !inputExists {
		msg := fmt.Sprintf("%s: -i INPUT: File does not exist.\n%s", vars.CommandFlag, inputValue)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	inputFileMode := utils.GetFileMode(inputValue)

	if inputFileMode == 0 {
		msg := fmt.Sprintf("%s: -i INPUT: File does not appear to be a valid encrypted file.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	var apiValue string
	var apiIsValidPath bool
	var apiExists bool
	var apiIsValidJson bool
	var password string

	isAPISet, apiHasValue := utils.IsFlagPassed("A")
	if isAPISet {
		apiValue, _ = utils.GetFlagValue("A")
		apiIsValidPath = utils.IsValidPath(apiValue)
		apiExists = utils.FileExists(apiValue)
		apiIsValidJson, _ = utils.ValidateJSON(apiValue)
	}

	_ = apiHasValue
	_ = apiExists
	_ = apiIsValidJson

	var lockValidationRequired bool

	switch inputFileMode {
	case 1: // classic - do nothing
		lockValidationRequired = false
	case 2: // classic locked - apikey and password required
		lockValidationRequired = true

	case 3: // hybrid - do nothing
		lockValidationRequired = false
	case 4: // hybrid locked - apikey and password required
		lockValidationRequired = true
	default:
		lockValidationRequired = false
		msg := fmt.Sprintf("%s: -i INPUT: Unable to determine encryption mode. Invalid File.\n%s\n", vars.CommandFlag, inputValue)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if isAPISet && lockValidationRequired {
		//

		if !isAPISet {
			msg := fmt.Sprintf("%s: LOCKED FILE: -A APIKEY: flag is required.\n", vars.CommandFlag)
			utils.HandleFailure(msg)
			os.Exit(1)
		}

		if isAPISet && !apiHasValue {
			msg := fmt.Sprintf("%s: LOCKED FILE: -A APIKEY: requires a value.\n", vars.CommandFlag)
			utils.HandleFailure(msg)
			os.Exit(1)
		}

		if isAPISet && !apiIsValidPath {
			msg := fmt.Sprintf("%s: LOCKED FILE: -A APIKEY: requires a valid file path.\n", vars.CommandFlag)
			utils.HandleFailure(msg)
			os.Exit(1)
		}

		password := utils.PromptPassword()

		_ = password

		//
	}

	_ = apiIsValidPath
	_ = password

	// -- Output

	isOutputSet, outputHasValue := utils.IsFlagPassed("o")
	outputValue, _ := utils.GetFlagValue("o")
	outputIsValidPath := utils.IsValidPath(outputValue)
	outputExists := utils.FileExists(outputValue)

	if isOutputSet && !outputHasValue {
		msg := fmt.Sprintf("%s: -o OUTPUT: requires a value.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	if isOutputSet && !outputIsValidPath {
		msg := fmt.Sprintf("%s: -o OUTPUT: requires a valid file path.\n", vars.CommandFlag)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	// confirm overwrite.....

	if isOutputSet && outputExists {
		utils.ConfirmOverwrite(outputValue)
	}

	// Get Private Key

	privateKey, err := utils.FetchFileKey(keyValue, "key")
	if err != nil {
		msg := fmt.Sprintf("%s: -k KEY: Error: %s", vars.CommandFlag, err)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	keyBytes, err := os.ReadFile(keyValue)
	if err != nil {
		msg := fmt.Sprintf("%s: -k KEY: Could not read key file.\n%s\n", vars.CommandFlag, err)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	var pKey syscrypt.PrivateKeyWrapper
	err = json.Unmarshal(keyBytes, &pKey)
	if err != nil {
		msg := fmt.Sprintf("%s: -k KEY: File is not a valid syscrypt key file. %v\n", vars.CommandFlag, err)
		utils.HandleFailure(msg)
		os.Exit(1)
	}

	DecryptFile(outputValue, inputValue, pKey)

	_ = err

}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func validateFile(inputValue string) (encryptionType bool, isLocked bool, isArmored bool) {

	//switch mode {
	//case 0x01:
	//	fmt.Printf("classic")
	//case 0x02:
	//	fmt.Printf("classic with lock")
	//case 0x03:
	//	fmt.Printf("post quantum")
	//case 0x04:
	//	fmt.Printf("post quantum with lock")
	//default:
	//	fmt.Printf("invalid")
	//}

	return false, false, false

	//file, err := os.Open(inputValue)
	//if err != nil {
	//	fmt.Printf("Error: cannot open %s\n", inputValue)
	//	os.Exit(1)
	//}
	//defer file.Close()

	//header := make([]byte, 1)
	//_, err = file.Read(header)
	//if err != nil {
	//	fmt.Println("Error: file is empty or unreadable")
	//	os.Exit(1)
	//}

	//file, err := os.Open(filePath)
	//if err != nil {
	//	return false, false, err
	//}
	//defer file.Close()

	//buffer := make([]byte, 512)
	//n, err := file.Read(buffer)
	//if err != nil && n == 0 {
	//	return false, false, err
	//}

	//preview := string(buffer[:n])
	//nPreview := strings.ReplaceAll(preview, vars.PrivateKeyHeader+"\n", "")
	//nPreview = strings.ReplaceAll(nPreview, vars.PrivateKeyFooter+"\n", "")
	//isBase64 := IsValidBase64WithLines(preview)

	//if strings.HasPrefix(preview, vars.PrivateKeyHeader) {
	//	isArmored = true
	//} else {
	//	isArmored = false
	//}

	//mode := header[0]

	//switch mode {
	//case 0x01:
	//	return 1 // Classical Encryption
	//case 0x02:
	//	return 2 // Classical Encryption [with lock]
	//case 0x03:
	//	return 3 // Post Quantum Encryption
	//case 0x04:
	//	return 4 // Post Quantum Encryption [with lock]
	//default:
	//	return 5 // Invalid or Armored
	//}

}

func DecryptFile(out string, input string, priv syscrypt.PrivateKeyWrapper) {

	outputValue, _ := utils.GetFlagValue("o")
	inputValue, _ := utils.GetFlagValue("i")

	isOutputSet, outputHasValue := utils.IsFlagPassed("o")
	isOutputValidPath := utils.IsValidPath(outputValue)
	outputExists := utils.FileExists(outputValue)
	passValue, _ := utils.GetFlagValue("P")
	apiValue, _ := utils.GetFlagValue("A")

	raw, _ := os.ReadFile(input)
	var blob []byte

	if strings.Contains(string(raw), vars.PrivateKeyHeader) {
		c := strings.ReplaceAll(string(raw), vars.PrivateKeyHeader, "")
		c = strings.ReplaceAll(c, vars.PrivateKeyFooter, "")
		c = strings.ReplaceAll(c, "\n", "")
		c = strings.ReplaceAll(c, "\r", "")
		blob, _ = base64.StdEncoding.DecodeString(strings.TrimSpace(c))
	} else {

		blob = raw

	}

	// unblind header

	ephPubX := blob[2:34]
	mode := blob[0] ^ ephPubX[0]
	serialSize := int(blob[1] ^ ephPubX[1])

	serialOffset := 34
	dataPtr := serialOffset + serialSize

	recoveredSerial := make([]byte, serialSize)
	for i := 0; i < serialSize; i++ {
		recoveredSerial[i] = blob[serialOffset+i] ^ ephPubX[(i+2)%32]
	}

	if priv.PrivateKey.Serial != string(recoveredSerial) {
		fmt.Printf("Access Denied: Serial Mismatch or invalid key type.\n")
		os.Exit(1)
	}

	isHybrid := (mode == 0x03 || mode == 0x04)
	isLocked := (mode == 0x02 || mode == 0x04)

	// Password Verification for Locked files
	var passKey []byte

	if isLocked {
		//
		apiKey, err := utils.FetchFileKey(apiValue, "key")
		if err != nil {
			fmt.Printf("Error: %s", err)
			os.Exit(1)
		}
		passKey = argon2.IDKey([]byte(passValue), []byte(apiKey), 1, 64*1024, 4, 32)

		vH := hkdf.New(sha256.New, passKey, nil, []byte("syscrypt-pass-verify"))
		vK := make([]byte, 32)
		io.ReadFull(vH, vK)
		vAead, _ := cc20.New(vK)
		res, err := vAead.Open(nil, make([]byte, 12), blob[dataPtr:dataPtr+vars.PassTagSize], nil)
		if err != nil || !bytes.Equal(res, []byte("SYSC-PASS-OK")) {
			fmt.Printf("Invalid Master Password.")
			os.Exit(1)
		}
		dataPtr += vars.PassTagSize
		//
	}

	// Slice Hybrid & Data

	var kyberCT []byte
	if isHybrid {
		kyberCT = blob[dataPtr : dataPtr+vars.KyberCTSize]
		dataPtr += vars.KyberCTSize
	}

	nonce := blob[dataPtr : dataPtr+12]
	ciphertext := blob[dataPtr+12:]

	// Final Key Reconstruction

	cleanPrivX := strings.TrimPrefix(priv.PrivateKey.Key, vars.PrivateKeyPrefixLabel)
	myPrivX, _ := hex.DecodeString(cleanPrivX)
	sharedX, _ := curve25519.X25519(myPrivX, ephPubX)

	var sharedML []byte
	if isHybrid {
		scheme := kyber768.Scheme()
		cleanML := strings.TrimPrefix(priv.PrivateKey.MLKEMKey, vars.PQPrivateKeyPrefixLabel)
		skBytes, _ := hex.DecodeString(cleanML)
		skK, _ := scheme.UnmarshalBinaryPrivateKey(skBytes)
		sharedML, _ = scheme.Decapsulate(skK, kyberCT)
	}

	combined := append(sharedX, sharedML...)
	if isLocked {
		combined = append(combined, passKey...)
	}

	h := hkdf.New(sha256.New, combined, nil, []byte("syscrypt-v1-hybrid"))
	symmK := make([]byte, 32)
	io.ReadFull(h, symmK)
	aead, _ := cc20.New(symmK)
	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		fmt.Printf("Decryption failed")
		os.Exit(1)
	}

	if isOutputSet {
		os.WriteFile(outputValue, plaintext, 0644)
		fmt.Printf("Successfully decrypted file to: %s\n", outputValue)
	}

	if !isOutputSet {
		fmt.Printf("%s\n", plaintext)
	}

	_ = inputValue
	_ = isOutputSet
	_ = outputHasValue
	_ = isOutputValidPath
	_ = outputExists
	_ = dataPtr
	_ = kyberCT
	_ = nonce
	_ = ciphertext
	_ = sharedX

	//outputValue, _ := utils.GetFlagValue("o")
	//inputValue, _ := utils.GetFlagValue("i")

	//isOutputSet, outputHasValue := utils.IsFlagPassed("o")
	//outputValue, _ := utils.GetFlagValue("k")
	//outputIsValidPath := utils.IsValidPath(outValue)
	//outputExists := utils.FileExists(outputValue)

	/*

		rawInput, err := os.ReadFile(inputValue)
		if err != nil {
			fmt.Printf("Error: Unable to read %s\n", inputValue)
			os.Exit(1)
		}

		var blob []byte
		inputStr := string(rawInput)

		if strings.Contains(inputStr, vars.PrivateKeyHeader) {
			content := strings.ReplaceAll(inputStr, vars.PrivateKeyHeader, "")
			content = strings.ReplaceAll(content, vars.PrivateKeyFooter, "")
			content = strings.ReplaceAll(content, "\n", "")
			content = strings.ReplaceAll(content, "\r", "")
			content = strings.TrimSpace(content)
			blob, err = base64.StdEncoding.DecodeString(content)
			if err != nil {
				fmt.Printf("Error: Malformed Armored Base64 data.\n")
				os.Exit(1)
			}

		} else {

			blob = rawInput

		}

		if len(blob) < 35 {
			fmt.Printf("Error: -i INPUT: File is too small to be a valid syscrypt encrypted message.\n")
			os.Exit(1)
		}

		mode := blob[0]
		serialSize := int(blob[1])

		ephOffset := 2 + serialSize
		headerOffset := ephOffset + 32

		if len(blob) < headerOffset {
			fmt.Printf("Error: Header truncated or invalid serial length.\n")
		}

		maskedSerial := blob[2:ephOffset]
		ephPubX := blob[ephOffset:headerOffset]

		recoveredBytes := make([]byte, serialSize)
		for i := 0; i < serialSize; i++ {
			recoveredBytes[i] = maskedSerial[i] ^ ephPubX[i%32]
		}
		fileSerial := string(recoveredBytes)

		if fileSerial != priv.PrivateKey.Serial {
			fmt.Printf("Error: File requires key [%s], but you provided [%s]\n", fileSerial, priv.PrivateKey.Serial)
			os.Exit(1)
		}

		var kyberCT, nonce, ciphertext []byte
		isHybrid := mode == 0x03 || mode == 0x04
		isLocked := mode == 0x02 || mode == 0x04

		if isHybrid {
			kyberCT = blob[headerOffset : headerOffset+1088]
			nonce = blob[headerOffset+1088 : headerOffset+1100]
			ciphertext = blob[headerOffset+1100:]
		} else {
			nonce = blob[headerOffset : headerOffset+12]
			ciphertext = blob[headerOffset+12:]
		}

		cleanPriv := strings.TrimPrefix(priv.PrivateKey.Key, vars.PrivateKeyPrefixLabel)
		myPrivX, _ := hex.DecodeString(cleanPriv)
		sharedX, _ := curve25519.X25519(myPrivX, ephPubX)

		var sharedML []byte
		if isHybrid {
			scheme := kyber768.Scheme()
			skBytes, _ := hex.DecodeString(priv.PrivateKey.MLKEMKey)
			skK, _ := scheme.UnmarshalBinaryPrivateKey(skBytes)
			sharedML, _ = scheme.Decapsulate(skK, kyberCT)
		}

		combined := append(sharedX, sharedML...)
		if isLocked {
			apiValue, _ := utils.GetFlagValue("A")
			passValue, _ := utils.GetFlagValue("P")

			passKey := argon2.IDKey([]byte(passValue), []byte(apiKey), 1, 64*1024, 4, 32)
			combined = append(combined, passKey...)
			defer utils.Zeroize(passKey)
		}

		h := hkdf.New(sha256.New, combined, nil, []byte("syscrypt-v1-hybrid"))
		symmKey := make([]byte, 32)
		io.ReadFull(h, symmKey)
		defer utils.Zeroize(symmKey)

		aead, _ := cc20.New(symmKey)
		plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
		if err != nil {
			fmt.Printf("Error: Decryption failed. Potentional data corruption or wrong password.")
			os.Exit(1)
		}

		if isOutputSet && outputIsValidPath {

			err = os.WriteFile(outputValue, plaintext, 0644)
			if err != nil {
				fmt.Printf("Error: Unable to write decrypted file %s\n", outputValue)
				os.Exit(1)
			}

		} else {
			/// do stuff that isnt output
		}
	*/

}