package main

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/asn1"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"log"
	"math/big"
	"net"
	"time"

	"go-ts/pkg/protocol"

	"filippo.io/edwards25519"
)

type ServerState struct {
	Step int
	A1   [16]byte
	A2   [100]byte

	Identity *ecdsa.PrivateKey

	LicensePriv  *edwards25519.Scalar
	LicensePub   *edwards25519.Point
	LicenseBlock []byte

	Alpha        []byte
	Beta         []byte
	SharedSecret []byte
	SharedIV     []byte
	SharedMac    [8]byte
}

func main() {
	addr, _ := net.ResolveUDPAddr("udp", ":9988")
	conn, _ := net.ListenUDP("udp", addr)
	log.Println("FakeServer listening on :9988")

	// 1. Generate Server Identity (P-256)
	privKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

	// 2. Generate Transport/License Key (Ed25519)
	var seed [32]byte
	rand.Read(seed[:])
	lPriv, _ := new(edwards25519.Scalar).SetBytesWithClamping(seed[:])
	lPub := new(edwards25519.Point).ScalarBaseMult(lPriv)

	// 3. Create 'l' (License Block) - 32 random bytes for now
	lData := make([]byte, 32)
	rand.Read(lData)

	state := &ServerState{
		Step:         0,
		Identity:     privKey,
		LicensePriv:  lPriv,
		LicensePub:   lPub,
		LicenseBlock: lData,
		Alpha:        make([]byte, 10),
		Beta:         make([]byte, 54),
	}
	rand.Read(state.Alpha)
	rand.Read(state.Beta)
	rand.Read(state.A2[:])

	buf := make([]byte, 4096)
	for {
		n, rAddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			continue
		}
		data := make([]byte, n)
		copy(data, buf[:n])
		handlePacket(conn, rAddr, data, state)
	}
}

func handlePacket(conn *net.UDPConn, addr *net.UDPAddr, data []byte, s *ServerState) {
	pkt, err := protocol.Decode(data, true)
	if err != nil {
		return
	}

	if pkt.Header.PacketType() == protocol.PacketTypeInit1 {
		if len(pkt.Data) > 5 && pkt.Data[4] == 0x00 { // Step 0
			log.Println("Recv Step 0, Sending Step 1")
			sendStep1(conn, addr, s)
		} else if len(pkt.Data) > 5 && pkt.Data[4] == 0x02 { // Step 2
			log.Println("Recv Step 2, Sending Step 3")
			sendStep3(conn, addr, s)
		}
	} else if pkt.Header.PacketType() == protocol.PacketTypeCommand {
		decrypted, err := decryptHandshake(pkt)
		if err == nil {
			sStr := string(decrypted)
			if len(sStr) > 8 && sStr[0:8] == "clientek" {
				log.Printf("Recv clientek (Decrypted): %s", sStr)
				if err := s.processClientEk(decrypted); err != nil {
					log.Printf("Error processing clientek: %v", err)
				} else {
					log.Println("Shared Secret Derived. Waiting for clientinit.")
				}
				sendAck(conn, addr, pkt.Header.PacketID)
				return
			} else if len(sStr) > 12 && sStr[0:12] == "clientinitiv" {
				log.Println("Recv clientinitiv. Sending initivexpand2...")
				sendInitivexpand2(conn, addr, s)
				return
			}
		}

		if len(s.SharedSecret) > 0 {
			s.decryptClientInit(pkt)
		} else if err != nil && pkt.Header.FlagUnencrypted() == false {
			// log.Printf("Decrypt failed: %v", err)
		}
	}
}

func decryptHandshake(pkt *protocol.Packet) ([]byte, error) {
	key := protocol.HandshakeKey
	nonce := protocol.HandshakeNonce

	meta := make([]byte, 5)
	binary.BigEndian.PutUint16(meta[0:2], pkt.Header.PacketID)
	binary.BigEndian.PutUint16(meta[2:4], pkt.Header.ClientID)
	meta[4] = pkt.Header.Type

	return protocol.DecryptEAX(key, nonce, meta, pkt.Data, pkt.Header.MAC[:])
}

func sendStep1(conn *net.UDPConn, addr *net.UDPAddr, s *ServerState) {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, int32(time.Now().Unix()))
	buf.WriteByte(0x01)
	rand.Read(s.A1[:])
	buf.Write(s.A1[:])

	pkt := protocol.NewPacket(protocol.PacketTypeInit1, buf.Bytes())
	pkt.Header.PacketID = 1
	copy(pkt.Header.MAC[:], []byte("TS3INIT1"))

	encoded, _ := pkt.Encode(false)
	conn.WriteToUDP(encoded, addr)
}

func sendStep3(conn *net.UDPConn, addr *net.UDPAddr, s *ServerState) {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, int32(time.Now().Unix()))
	buf.WriteByte(0x03)

	x := make([]byte, 64)
	rand.Read(x)
	n := make([]byte, 64)
	rand.Read(n)
	buf.Write(x)
	buf.Write(n)
	binary.Write(buf, binary.BigEndian, uint32(0))
	buf.Write(s.A2[:])

	pkt := protocol.NewPacket(protocol.PacketTypeInit1, buf.Bytes())
	pkt.Header.PacketID = 2
	copy(pkt.Header.MAC[:], []byte("TS3INIT1"))

	encoded, _ := pkt.Encode(false)
	conn.WriteToUDP(encoded, addr)
}

func sendInitivexpand2(conn *net.UDPConn, addr *net.UDPAddr, s *ServerState) {
	lStr := base64.StdEncoding.EncodeToString(s.LicenseBlock)
	base64.StdEncoding.EncodeToString(s.Beta) // unused?
	betaStr := base64.StdEncoding.EncodeToString(s.Beta)

	// Encode Omega (P-256 Public Key)
	pub := s.Identity.PublicKey
	omegaBytes := elliptic.Marshal(pub.Curve, pub.X, pub.Y)
	omegaStr := base64.StdEncoding.EncodeToString(omegaBytes)

	// Create Proof: Sign(SHA256(lBytes)) with Identity
	hash := sha256.Sum256(s.LicenseBlock)
	r, sb, err := ecdsa.Sign(rand.Reader, s.Identity, hash[:])
	if err != nil {
		log.Printf("Signing failed: %v", err)
	}

	type ECDSASignature struct {
		R, S *big.Int
	}
	sig := ECDSASignature{R: r, S: sb}
	proofBytes, _ := asn1.Marshal(sig)
	proofStr := base64.StdEncoding.EncodeToString(proofBytes)

	// Send command
	cmd := fmt.Sprintf("initivexpand2 l=%s beta=%s omega=%s ot=1 proof=%s tvd=C",
		lStr, betaStr, omegaStr, proofStr)

	pkt := protocol.NewPacket(protocol.PacketTypeCommand, []byte(cmd))
	pkt.Header.PacketID = 3

	// Encrypt with HandshakeKey
	key := protocol.HandshakeKey
	nonce := protocol.HandshakeNonce

	// Meta S->C (3 bytes: PID(2)+Type(1))
	meta := make([]byte, 3)
	binary.BigEndian.PutUint16(meta[0:2], pkt.Header.PacketID)
	meta[2] = pkt.Header.Type

	encData, mac, _ := protocol.EncryptEAX(key, nonce, meta, pkt.Data)
	pkt.Data = encData
	copy(pkt.Header.MAC[:], mac)

	encoded, _ := pkt.Encode(false)
	conn.WriteToUDP(encoded, addr)
}

func sendAck(conn *net.UDPConn, addr *net.UDPAddr, pid uint16) {
	pkt := protocol.NewPacket(protocol.PacketTypeAck, nil)
	pkt.Header.PacketID = pid

	// Encrypt ACK
	key := protocol.HandshakeKey
	nonce := protocol.HandshakeNonce
	meta := make([]byte, 3)
	binary.BigEndian.PutUint16(meta[0:2], pid)
	meta[2] = pkt.Header.Type

	encData, mac, _ := protocol.EncryptEAX(key, nonce, meta, pkt.Data)
	pkt.Data = encData
	copy(pkt.Header.MAC[:], mac)

	encoded, _ := pkt.Encode(false)
	conn.WriteToUDP(encoded, addr)
}

func (s *ServerState) deriveSharedSecret(clientEkPub []byte) error {
	clientPoint, err := new(edwards25519.Point).SetBytes(clientEkPub)
	if err != nil {
		return fmt.Errorf("invalid client point: %v", err)
	}

	// Server negates client point
	clientPoint.Negate(clientPoint)

	sharedPoint := new(edwards25519.Point).ScalarMult(s.LicensePriv, clientPoint)
	sharedBytes := sharedPoint.Bytes()

	sharedBytes[31] ^= 0x80

	hash := sha512.Sum512(sharedBytes)
	s.SharedSecret = hash[:]

	s.SharedIV = make([]byte, 64)
	copy(s.SharedIV, s.SharedSecret)
	for i := 0; i < 10; i++ {
		s.SharedIV[i] ^= s.Alpha[i]
	}
	if len(s.Beta) >= 54 {
		for i := 0; i < 54; i++ {
			s.SharedIV[10+i] ^= s.Beta[i]
		}
	}

	macHash := sha1.Sum(s.SharedIV)
	copy(s.SharedMac[:], macHash[0:8])

	log.Printf("Shared Secret Derived! IV: %s", hex.EncodeToString(s.SharedIV))
	return nil
}

func (s *ServerState) processClientEk(data []byte) error {
	str := string(data)
	start := "ek="
	idx := 0
	for i := 0; i < len(str)-len(start); i++ {
		if str[i:i+3] == start {
			idx = i + 3
			break
		}
	}
	if idx == 0 {
		return fmt.Errorf("no ek found")
	}

	end := idx
	for i := idx; i < len(str); i++ {
		if str[i] == ' ' {
			end = i
			break
		}
	}

	ekStr := str[idx:end]
	ekBytes, err := base64.StdEncoding.DecodeString(ekStr)
	if err != nil {
		return err
	}

	return s.deriveSharedSecret(ekBytes)
}

func (s *ServerState) decryptClientInit(pkt *protocol.Packet) {
	crypto := &protocol.CryptoState{
		SharedIV:     s.SharedIV,
		SharedMac:    s.SharedMac[:],
		GenerationID: 0,
	}

	key, nonce := crypto.GenerateKeyNonce(&pkt.Header, true) // isClientToServer

	meta := make([]byte, 5)
	binary.BigEndian.PutUint16(meta[0:2], pkt.Header.PacketID)
	binary.BigEndian.PutUint16(meta[2:4], pkt.Header.ClientID)
	meta[4] = pkt.Header.Type

	dec, err := protocol.DecryptEAX(key, nonce, meta, pkt.Data, pkt.Header.MAC[:])
	if err != nil {
		log.Printf("Decryption Failed: %v", err)
		return
	}

	log.Printf(">>> DECRYPTED CLIENTINIT <<<\n%s\n", string(dec))
}