package client

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

	// Unused but placeholder
	"go-ts/pkg/protocol"
	"go-ts/pkg/transport"

	"filippo.io/edwards25519"
)

// HandshakeState tracks the progress of the connection initialization.
type HandshakeState struct {
	Step int

	// Init Cookies/Buffers
	A0 [4]byte
	A1 [16]byte
	A2 [100]byte

	// Puzzle Data
	X     *big.Int
	N     *big.Int
	Level uint32

	// Identity
	IdentityKey *ecdsa.PrivateKey
	Alpha       []byte // 10 random bytes

	// Server Data
	Beta    []byte
	Omega   []byte // Server Public Key (DER)
	License []byte

	// Crypto
	ClientEkPub  [32]byte
	ClientEkPriv [32]byte
	ClientScalar *edwards25519.Scalar // Client ephemeral private key (scalar)

	SharedSecret []byte
	SharedIV     []byte
	SharedMac    []byte

	IdentityOffset uint64 // Extracted/Mined offset
	IdentityLevel  int

	Conn *transport.TS3Conn
}

func NewHandshakeState(conn *transport.TS3Conn) (*HandshakeState, error) {
	// Generate Identity On Startup (or load from disk in future)

	privKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate identity key: %v", err)
	}

	alpha := make([]byte, 10)
	if _, err := rand.Read(alpha); err != nil {
		return nil, fmt.Errorf("failed to generate alpha: %v", err)
	}

	// Generate Client Ephemeral Key (TS3 Style)
	// TS3 uses a random 32-byte scalar with the high bit masked (&= 0x7F).
	// It does NOT use standard Ed25519 clamping (bits 0-2 cleared, bit 254 set).
	var clientScalar *edwards25519.Scalar
	var clientKeyBytes [32]byte

	for {
		if _, err := rand.Read(clientKeyBytes[:]); err != nil {
			return nil, err
		}
		clientKeyBytes[31] &= 0x7F // Mask high bit (positive scalar)

		// Try to load as scalar. Might fail if >= L (very unlikely)
		clientScalar, err = new(edwards25519.Scalar).SetCanonicalBytes(clientKeyBytes[:])
		if err == nil {
			break
		}
	}

	return &HandshakeState{
		Step:         0,
		Conn:         conn,
		IdentityKey:  privKey,
		Alpha:        alpha,
		ClientScalar: clientScalar,
	}, nil
}

// ... SendPacket0, HandlePacket1, SendPacket2, HandlePacket3 logic unchanged ...
// (Omitting for brevity if replacing, but user asked for full file usually.
// Just including necessary parts and new methods)

func (h *HandshakeState) SendPacket0() error {
	buf := new(bytes.Buffer)
	ts := int32(time.Now().Unix()) - 1356998400
	binary.Write(buf, binary.BigEndian, ts)
	buf.WriteByte(0x00) // Step 0
	now := int32(time.Now().Unix())
	binary.Write(buf, binary.BigEndian, now)
	rand.Read(h.A0[:])
	buf.Write(h.A0[:])
	buf.Write(make([]byte, 8)) // 8 Zeros
	pkt := protocol.NewPacket(protocol.PacketTypeInit1, buf.Bytes())
	pkt.Header.PacketID = 101
	pkt.Header.ClientID = 0
	pkt.Header.MAC = protocol.HandshakeMac
	return h.Conn.SendPacket(pkt)
}

func (h *HandshakeState) HandlePacket1(pkt *protocol.Packet) error {
	if len(pkt.Data) < 21 {
		return fmt.Errorf("packet 1 too short")
	}
	if pkt.Data[0] != 0x01 {
		return fmt.Errorf("expected step 1, got %d", pkt.Data[0])
	}
	copy(h.A1[:], pkt.Data[1:17])
	h.Step = 1
	return nil
}

func (h *HandshakeState) SendPacket2() error {
	buf := new(bytes.Buffer)
	ts := int32(time.Now().Unix()) - 1356998400
	binary.Write(buf, binary.BigEndian, ts)
	buf.WriteByte(0x02) // Step 2
	buf.Write(h.A1[:])
	a0Rev := [4]byte{h.A0[3], h.A0[2], h.A0[1], h.A0[0]}
	buf.Write(a0Rev[:])
	pkt := protocol.NewPacket(protocol.PacketTypeInit1, buf.Bytes())
	pkt.Header.PacketID = 102
	pkt.Header.MAC = protocol.HandshakeMac
	return h.Conn.SendPacket(pkt)
}

func (h *HandshakeState) HandlePacket3(pkt *protocol.Packet) error {
	if len(pkt.Data) < 233 {
		return fmt.Errorf("packet 3 too short")
	}
	if pkt.Data[0] != 0x03 {
		return fmt.Errorf("expected step 3, got %d", pkt.Data[0])
	}
	h.X = new(big.Int).SetBytes(pkt.Data[1:65])
	h.N = new(big.Int).SetBytes(pkt.Data[65:129])
	h.Level = binary.BigEndian.Uint32(pkt.Data[129:133])
	copy(h.A2[:], pkt.Data[133:233])
	h.Step = 3
	log.Printf("Received Puzzle: Level=%d", h.Level)
	return h.SendPacket4()
}

func (h *HandshakeState) SendPacket4() error {
	e := new(big.Int).Lsh(big.NewInt(1), uint(h.Level))
	y := new(big.Int).Exp(h.X, e, h.N)
	yBytes := y.Bytes()
	yPadded := make([]byte, 64)
	if len(yBytes) > 64 {
		copy(yPadded, yBytes[len(yBytes)-64:])
	} else {
		copy(yPadded[64-len(yBytes):], yBytes)
	}
	buf := new(bytes.Buffer)
	ts := int32(time.Now().Unix()) - 1356998400
	binary.Write(buf, binary.BigEndian, ts)
	buf.WriteByte(0x04) // Step 4
	xPadded := make([]byte, 64)
	xb := h.X.Bytes()
	if len(xb) > 64 {
		copy(xPadded, xb[len(xb)-64:])
	} else {
		copy(xPadded[64-len(xb):], xb)
	}
	buf.Write(xPadded)
	nPadded := make([]byte, 64)
	nb := h.N.Bytes()
	if len(nb) > 64 {
		copy(nPadded, nb[len(nb)-64:])
	} else {
		copy(nPadded[64-len(nb):], nb)
	}
	buf.Write(nPadded)
	binary.Write(buf, binary.BigEndian, h.Level)
	buf.Write(h.A2[:])
	buf.Write(yPadded)
	cmdData := h.GenerateClientInitIV()
	buf.Write([]byte(cmdData))
	pkt := protocol.NewPacket(protocol.PacketTypeInit1, buf.Bytes())
	pkt.Header.PacketID = 103
	pkt.Header.MAC = protocol.HandshakeMac
	log.Println("Sending Packet 4 with Solution and clientinitiv...")
	return h.Conn.SendPacket(pkt)
}

// ProcessInitivexpand2 handles the decrypted command logic
func (h *HandshakeState) ProcessInitivexpand2(cmdArgs map[string]string) error {
	if h.Step >= 5 {
		log.Println("Ignoring duplicate initivexpand2 (Step already advanced)")
		return nil
	}

	lStr, ok := cmdArgs["l"]
	if !ok {
		return errors.New("missing license (l)")
	}

	lData, err := base64.StdEncoding.DecodeString(lStr)
	if err != nil {
		return err
	}

	betaStr, ok := cmdArgs["beta"]
	if !ok {
		return errors.New("missing beta")
	}
	h.Beta, err = base64.StdEncoding.DecodeString(betaStr)
	if err != nil {
		return err
	}

	// 1. Derive Server Public Key (Edwards Y)
	serverEdPubBytes, err := protocol.ParseLicenseAndDeriveKey(lData)
	if err != nil {
		return fmt.Errorf("LICENSE FAIL: %v", err)
	}

	// Load Server Public Key as Edwards Point
	serverPoint, err := new(edwards25519.Point).SetBytes(serverEdPubBytes)
	if err != nil {
		return fmt.Errorf("invalid server public key point: %v", err)
	}

	// 2. Client Ephemeral Key (Scalar) is pre-generated in NewHandshakeState (h.ClientScalar)
	// Compute Client Public Key (Point) = Scalar * Base
	clientPubPoint := new(edwards25519.Point).ScalarBaseMult(h.ClientScalar)
	h.ClientEkPub = [32]byte(clientPubPoint.Bytes())

	// 3. Calculate Shared Secret (Ed25519 Scalar Mult)

	// Negate Server Public Key (TS3/Punisher.NaCl logic)
	serverPointNeg := new(edwards25519.Point).Negate(serverPoint)

	// Multiply: Result = Scalar * (-ServerPoint)
	sharedPoint := new(edwards25519.Point).ScalarMult(h.ClientScalar, serverPointNeg)
	sharedBytes := sharedPoint.Bytes()

	// XOR the last byte with 0x80 (Flip sign bit of X coordinate)
	sharedBytes[31] ^= 0x80

	// 4. SHA512 Hash of the result
	hash := sha512.Sum512(sharedBytes)
	h.SharedSecret = hash[:]

	h.SharedIV = make([]byte, 64)
	copy(h.SharedIV, h.SharedSecret)

	// XOR operations
	// SharedIV[0..10] xor alpha
	// SharedIV[10..64] xor beta

	// Alpha is 10 bytes (h.Alpha)
	for i := 0; i < 10; i++ {
		h.SharedIV[i] ^= h.Alpha[i]
	}

	// Beta should be 54 bytes
	log.Printf("Debug - Beta Length: %d", len(h.Beta))
	if len(h.Beta) >= 54 {
		for i := 0; i < 54; i++ {
			h.SharedIV[10+i] ^= h.Beta[i]
		}
	}

	// SharedMac = SHA1(SharedIV)[0..8]
	macHash := sha1.Sum(h.SharedIV)
	h.SharedMac = make([]byte, 8)
	copy(h.SharedMac, macHash[0:8])

	log.Printf("Debug - SharedSecret (SHA512): %s", hex.EncodeToString(h.SharedSecret))
	log.Printf("Debug - SharedIV: %s", hex.EncodeToString(h.SharedIV))
	log.Printf("Debug - SharedMac: %s", hex.EncodeToString(h.SharedMac[:]))
	log.Println("Shared Secret & Keys Calculated using TS3 Ed25519 logic.")

	return h.SendClientEk()
}

func (h *HandshakeState) SendClientEk() error {
	// clientek ek={ek} proof={proof}

	// ek = base64(client_public_key) [Ed25519 Compressed Point]
	ekStr := base64.StdEncoding.EncodeToString(h.ClientEkPub[:])

	// proof = base64(ecdsa_sign(client_public_key + beta))
	// "sign must be done with the private key from the identity keypair." (P-256)

	proofBuf := append(h.ClientEkPub[:], h.Beta...)
	hash := sha256.Sum256(proofBuf)

	r, s, err := ecdsa.Sign(rand.Reader, h.IdentityKey, hash[:])
	if err != nil {
		return err
	}

	// Encode Signature (ASN.1 DER)
	// Reverting to DER as server uses DER.

	sigContent := new(bytes.Buffer)
	writeBigInt(sigContent, r)
	writeBigInt(sigContent, s)

	sigSeq := new(bytes.Buffer)
	sigSeq.WriteByte(0x30)
	writeLength(sigSeq, sigContent.Len())
	sigSeq.Write(sigContent.Bytes())

	proofBytes := sigSeq.Bytes()
	proofStr := base64.StdEncoding.EncodeToString(proofBytes)

	log.Printf("Debug - ClientEk: %s", ekStr)
	log.Printf("Debug - Proof (DER): %s", hex.EncodeToString(proofBytes))

	// Construct Command
	cmd := fmt.Sprintf("clientek ek=%s proof=%s", protocol.Escape(ekStr), protocol.Escape(proofStr))

	// Packet 1 (New counter logic? Spec: "clientek already has the packet id 1")
	// This is the START of the new encrypted session?
	// "The normal packet id counting starts with this packet."

	buf := new(bytes.Buffer)
	buf.Write([]byte(cmd))

	pkt := protocol.NewPacket(protocol.PacketTypeCommand, buf.Bytes())
	pkt.Header.PacketID = 1                        // Reset to 1
	pkt.Header.SetType(protocol.PacketTypeCommand) // Ensure flag (NewProtocol? Unencrypted?)

	// Encryption?
	// "All Command ... Packets must get encrypted." (Using OLD Handshake keys? Or NEW?)
	// "The crypto handshake is now completed. The normal encryption scheme ... is from now on used."
	// Usually implies clientek IS encrypted with the NEW keys.

	// Add PacketFlagNewProtocol
	pkt.Header.Type |= protocol.PacketFlagNewProtocol // 0x20

	// Encryption
	// Try using HandshakeKey like initivexpand2 instead of SharedSecret
	// The crypto switch might happen AFTER clientek is accepted
	key := protocol.HandshakeKey
	nonce := protocol.HandshakeNonce

	// Prepare Meta for EAX
	// Meta for Client->Server: PID(2) + CID(2) + PT(1) = 5 bytes
	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

	// Encrypt
	cipherDict, mac, err := protocol.EncryptEAX(key, nonce, meta, pkt.Data)
	if err != nil {
		return err
	}

	pkt.Data = cipherDict
	copy(pkt.Header.MAC[:], mac)

	log.Println("Sending clientek (Packet 1) [Encrypted]")
	h.Step = 5
	return h.Conn.SendPacket(pkt)
}

func (h *HandshakeState) GenerateClientInitIV() string {
	// ... (Existing implementation) ...
	// Copy from previous step

	alphaStr := base64.StdEncoding.EncodeToString(h.Alpha)
	omegaStr := h.GenerateOmega()

	ip := "127.0.0.1"
	if addr, ok := h.Conn.RemoteAddr().(*net.UDPAddr); ok {
		ip = addr.IP.String()
	} else if h.Conn.RemoteAddr() != nil {
		ip = h.Conn.RemoteAddr().String()
		if host, _, err := net.SplitHostPort(ip); err == nil {
			ip = host
		}
	}

	return fmt.Sprintf("clientinitiv alpha=%s omega=%s ot=1 ip=%s",
		protocol.Escape(alphaStr), protocol.Escape(omegaStr), protocol.Escape(ip))
}

func (h *HandshakeState) GenerateOmega() string {
	// ... (Existing implementation) ...
	pub := h.IdentityKey.PublicKey

	content := new(bytes.Buffer)
	content.Write([]byte{0x03, 0x02, 0x07, 0x00})
	content.Write([]byte{0x02, 0x01, 0x20})
	writeBigInt(content, pub.X)
	writeBigInt(content, pub.Y)
	seq := new(bytes.Buffer)
	seq.WriteByte(0x30)
	writeLength(seq, content.Len())
	seq.Write(content.Bytes())
	return base64.StdEncoding.EncodeToString(seq.Bytes())
}

func writeBigInt(buf *bytes.Buffer, n *big.Int) {
	b := n.Bytes()
	buf.WriteByte(0x02)
	padded := b
	if len(b) > 0 && b[0] >= 0x80 {
		padded = make([]byte, len(b)+1)
		padded[0] = 0x00
		copy(padded[1:], b)
	} else if len(b) == 0 {
		padded = []byte{0x00}
	}
	writeLength(buf, len(padded))
	buf.Write(padded)
}

func writeLength(buf *bytes.Buffer, length int) {
	if length < 128 {
		buf.WriteByte(byte(length))
	} else {
		s := fmt.Sprintf("%x", length)
		if len(s)%2 != 0 {
			s = "0" + s
		}
		b, _ := hex.DecodeString(s)
		buf.WriteByte(0x80 | byte(len(b)))
		buf.Write(b)
	}
}

// ImproveSecurityLevel mines a counter to achieve the target security level.
func (h *HandshakeState) ImproveSecurityLevel(targetLevel int) {
	omega := h.GenerateOmega() // Base64 of ASN.1 Public Key

	// Start from current offset (usually 0)
	counter := h.IdentityOffset

	log.Printf("Mining Identity Level %d... ", targetLevel)

	for {
		// Construct data: Omega + Counter (ASCII)
		data := fmt.Sprintf("%s%d", omega, counter)

		// SHA1
		hash := sha1.Sum([]byte(data))

		// Count leading zero bits
		zeros := countLeadingZeros(hash[:])

		if zeros >= targetLevel {
			h.IdentityLevel = zeros
			h.IdentityOffset = counter
			log.Printf("Found! Offset=%d, Level=%d\n", counter, zeros)
			return
		}

		counter++
		if counter%100000 == 0 {
			// fmt.Print(".")
		}
	}
}

func countLeadingZeros(hash []byte) int {
	zeros := 0
	for _, b := range hash {
		if b == 0 {
			zeros += 8
		} else {
			// Count bits in this byte
			for i := 7; i >= 0; i-- {
				if (b>>i)&1 == 0 {
					zeros++
				} else {
					return zeros
				}
			}
			return zeros
		}
	}
	return zeros
}