package widevine

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"math/big"
	"strings"

	"github.com/golang/protobuf/proto"
)

// Constants
const (
	DeviceName = "chrome_1610"
)

// CDM represents the Content Decryption Module
type CDM struct {
	Devices  map[string]*CDMDevice
	Sessions map[string]*Session
}

// NewCDM creates a new CDM instance
func NewCDM() *CDM {
	return &CDM{
		Devices: map[string]*CDMDevice{
			DeviceName: NewCDMDevice(DeviceName, nil, nil, nil),
		},
		Sessions: make(map[string]*Session),
	}
}

// CheckPSSH validates and processes PSSH data
func (cdm *CDM) CheckPSSH(psshB64 string) ([]byte, error) {
	systemID := []byte{237, 239, 139, 169, 121, 214, 74, 206, 163, 200, 39, 220, 213, 29, 33, 237}

	// Pad the base64 string if needed
	if len(psshB64)%4 != 0 {
		psshB64 += strings.Repeat("=", 4-(len(psshB64)%4))
	}

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

	if len(pssh) < 30 {
		return pssh, nil
	}

	if !bytes.Equal(pssh[12:28], systemID) {
		newPssh := []byte{0, 0, 0}
		newPssh = append(newPssh, byte(32+len(pssh)))
		newPssh = append(newPssh, []byte("pssh")...)
		newPssh = append(newPssh, []byte{0, 0, 0, 0}...)
		newPssh = append(newPssh, systemID...)
		newPssh = append(newPssh, []byte{0, 0, 0, 0}...)
		newPssh = append(newPssh, byte(len(pssh)))
		newPssh = append(newPssh, pssh...)
		return newPssh, nil
	}

	return pssh, nil
}

// OpenSession opens a new DRM session
func (cdm *CDM) OpenSession(initDataB64, deviceName string, offline, raw bool) (string, error) {
	initData, err := cdm.CheckPSSH(initDataB64)
	if err != nil {
		return "", err
	}

	device, ok := cdm.Devices[deviceName]
	if !ok {
		return "", errors.New("device not found")
	}

	var sessionID []byte
	if device.IsAndroid {
		randHex := ""
		choices := "ABCDEF0123456789"
		for i := 0; i < 16; i++ {
			randHex += string(choices[randInt(len(choices))])
		}
		counter := "01"
		rest := "00000000000000"
		sessionID = []byte(randHex + counter + rest)
	} else {
		sessionID = make([]byte, 16)
		if _, err := rand.Read(sessionID); err != nil {
			return "", err
		}
	}

	var session *Session
	parsedInitData, err := cdm.ParseInitData(initData)
	if err == nil {
		session = NewSession(sessionID, parsedInitData, device, offline)
	} else if raw {
		session = NewSession(sessionID, initData, device, offline)
	} else {
		return "", errors.New("unable to parse init data")
	}

	sessionIDHex := hex.EncodeToString(sessionID)
	cdm.Sessions[sessionIDHex] = session

	return sessionIDHex, nil
}

// ParseInitData parses the initialization data
func (cdm *CDM) ParseInitData(initData []byte) (*WidevineCencHeader, error) {
	var cencHeader WidevineCencHeader

	if len(initData) < 32 {
		return nil, errors.New("init data too short")
	}

	err := proto.Unmarshal(initData[32:], &cencHeader)
	if err == nil {
		return &cencHeader, nil
	}

	// Try HBO Max format
	psshBox, err := ParsePSSHBox(initData)
	if err != nil {
		return nil, errors.New("unable to parse init data format")
	}

	err = proto.Unmarshal(psshBox.Data, &cencHeader)
	if err != nil {
		return nil, errors.New("unable to parse init data format")
	}

	return &cencHeader, nil
}

// CloseSession closes a DRM session
func (cdm *CDM) CloseSession(sessionID string) bool {
	if _, exists := cdm.Sessions[sessionID]; exists {
		delete(cdm.Sessions, sessionID)
		return true
	}
	return false
}

// SetServiceCertificate sets the service certificate for a session
func (cdm *CDM) SetServiceCertificate(sessionID string, certData []byte) bool {
	session, exists := cdm.Sessions[sessionID]
	if !exists {
		return false
	}

	var signedMessage SignedMessage
	if err := proto.Unmarshal(certData, &signedMessage); err == nil {
		var serviceCert SignedDeviceCertificate
		if err := proto.Unmarshal(signedMessage.Msg, &serviceCert); err == nil {
			session.ServiceCertificate = &serviceCert
			session.PrivacyMode = true
			return true
		}
	}

	var serviceCert SignedDeviceCertificate
	if err := proto.Unmarshal(certData, &serviceCert); err != nil {
		return false
	}

	session.ServiceCertificate = &serviceCert
	session.PrivacyMode = true
	return true
}

// GetLicenseRequest generates a license request
func (cdm *CDM) GetLicenseRequest(sessionID string) ([]byte, error) {
	session, exists := cdm.Sessions[sessionID]
	if !exists {
		return nil, errors.New("session ID doesn't exist")
	}

	var licenseRequest proto.Message

	switch initData := session.InitData.(type) {
	case *WidevineCencHeader:
		licenseRequest = &SignedLicenseRequest{
			Type: SignedLicenseRequest_LICENSE_REQUEST.Enum(),
			Msg: &LicenseRequest{
				Type: LicenseRequest_NEW.Enum(),
				KeyControlNonce: proto.Uint32(1093602366),
				ProtocolVersion: ProtocolVersion_CURRENT.Enum(),
				ContentId: &LicenseRequest_ContentIdentification{
					CencId: &LicenseRequest_ContentIdentification_Cenc{
						LicenseType: func() *LicenseType {
							if session.Offline {
								return LicenseType_OFFLINE.Enum()
							}
							return LicenseType_DEFAULT.Enum()
						}(),
						RequestId: session.SessionId,
						Pssh:      initData,
					},
				},
			},
		}
	case []byte:
		licenseRequest = &SignedLicenseRequestRaw{
			Type: SignedLicenseRequestRaw_LICENSE_REQUEST.Enum(),
			Msg: &LicenseRequestRaw{
				Type: LicenseRequestRaw_NEW.Enum(),
				KeyControlNonce: proto.Uint32(1093602366),
				ProtocolVersion: ProtocolVersion_CURRENT.Enum(),
				ContentId: &LicenseRequestRaw_ContentIdentification{
					CencId: &LicenseRequestRaw_ContentIdentification_Cenc{
						LicenseType: func() *LicenseType {
							if session.Offline {
								return LicenseType_OFFLINE.Enum()
							}
							return LicenseType_DEFAULT.Enum()
						}(),
						RequestId: session.SessionId,
						Pssh:      initData,
					},
				},
			},
		}
	default:
		return nil, errors.New("unsupported init data type")
	}

	if session.PrivacyMode && session.ServiceCertificate != nil {
		encryptedClientID, err := cdm.encryptClientID(session)
		if err != nil {
			return nil, err
		}

		switch lr := licenseRequest.(type) {
		case *SignedLicenseRequest:
			lr.Msg.EncryptedClientId = encryptedClientID
		case *SignedLicenseRequestRaw:
			lr.Msg.EncryptedClientId = encryptedClientID
		}
	} else {
		switch lr := licenseRequest.(type) {
		case *SignedLicenseRequest:
			lr.Msg.ClientId = session.Device.ClientID
		case *SignedLicenseRequestRaw:
			lr.Msg.ClientId = session.Device.ClientID
		}
	}

	// Serialize the message to sign it
	var msgBytes []byte
	var err error

	switch lr := licenseRequest.(type) {
	case *SignedLicenseRequest:
		msgBytes, err = proto.Marshal(lr.Msg)
	case *SignedLicenseRequestRaw:
		msgBytes, err = proto.Marshal(lr.Msg)
	}
	if err != nil {
		return nil, err
	}

	session.LicenseRequest = msgBytes

	// Sign the license request
	signature, err := session.Device.Sign(msgBytes)
	if err != nil {
		return nil, err
	}

	switch lr := licenseRequest.(type) {
	case *SignedLicenseRequest:
		lr.Signature = signature
	case *SignedLicenseRequestRaw:
		lr.Signature = signature
	}

	// Serialize the complete license request
	requestBytes, err := proto.Marshal(licenseRequest)
	if err != nil {
		return nil, err
	}

	return requestBytes, nil
}

// encryptClientID encrypts the client ID for privacy mode
func (cdm *CDM) encryptClientID(session *Session) (*EncryptedClientIdentification, error) {
	clientIDBytes, err := proto.Marshal(session.Device.ClientID)
	if err != nil {
		return nil, err
	}

	// Add PKCS7 padding
	blockSize := 16
	padding := blockSize - (len(clientIDBytes) % blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	clientIDBytes = append(clientIDBytes, padText...)

	// Generate AES key and IV
	aesKey := make([]byte, 16)
	if _, err := rand.Read(aesKey); err != nil {
		return nil, err
	}

	iv := make([]byte, aes.BlockSize)
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	// Encrypt the client ID
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	encryptedClientID := make([]byte, len(clientIDBytes))
	mode.CryptBlocks(encryptedClientID, clientIDBytes)

	// Encrypt the AES key with RSA
	pubKey, err := parseRSAPublicKey(session.ServiceCertificate.DeviceCertificate.PublicKey)
	if err != nil {
		return nil, err
	}

	encryptedKey, err := rsa.EncryptOAEP(sha1.New(), rand.Reader, pubKey, aesKey, nil)
	if err != nil {
		return nil, err
	}

	return &EncryptedClientIdentification{
		ServiceId:                   string(session.ServiceCertificate.DeviceCertificate.ServiceId),
		ServiceCertificateSerialNumber: session.ServiceCertificate.DeviceCertificate.SerialNumber,
		EncryptedClientId:           encryptedClientID,
		EncryptedClientIdIv:         iv,
		EncryptedPrivacyKey:         encryptedKey,
	}, nil
}

// parseRSAPublicKey parses an RSA public key from bytes
func parseRSAPublicKey(keyBytes []byte) (*rsa.PublicKey, error) {
	// This is a simplified version - actual implementation depends on the key format
	// In a real implementation, you would properly parse the ASN.1 encoded key
	if len(keyBytes) < 12 {
		return nil, errors.New("invalid key length")
	}

	// Assuming the key is in a simple format for this example
	// In reality, you'd use x509.ParsePKIXPublicKey or similar
	n := new(big.Int)
	n.SetBytes(keyBytes[8:]) // Skip some header bytes
	return &rsa.PublicKey{
		N: n,
		E: 65537, // Common exponent
	}, nil
}

// ProvideLicense processes the license response
func (cdm *CDM) ProvideLicense(sessionID string, license []byte) error {
	session, exists := cdm.Sessions[sessionID]
	if !exists {
		return errors.New("session ID doesn't exist")
	}

	if session.LicenseRequest == nil {
		return errors.New("generate a license request first")
	}

	var signedLicense SignedLicense
	if err := proto.Unmarshal(license, &signedLicense); err != nil {
		return errors.New("unable to parse license")
	}

	session.License = &signedLicense

	// Decrypt the session key
	sessionKey, err := session.Device.Decrypt(session.License.SessionKey)
	if err != nil {
		return errors.New("unable to decrypt session key")
	}

	if len(sessionKey) != 16 {
		return errors.New("invalid session key length")
	}

	session.SessionKey = sessionKey

	// Derive keys
	session.DerivedKeys = cdm.DeriveKeys(session.LicenseRequest, session.SessionKey)

	// Verify license signature
	licenseBytes, err := proto.Marshal(signedLicense.Msg)
	if err != nil {
		return err
	}

	hmacHash := getHMACSHA256Digest(licenseBytes, session.DerivedKeys.Auth1)
	if !hmac.Equal(hmacHash, signedLicense.Signature) {
		return errors.New("license signature mismatch")
	}

	// Decrypt content keys
	for _, key := range signedLicense.Msg.Keys {
		if key.GetType() == License_KeyContainer_SIGNING {
			continue
		}

		keyID := key.Id
		if keyID == nil {
			keyID = []byte(key.GetType().String())
		}

		decryptedKey, err := decryptKey(key.Key, key.Iv, session.DerivedKeys.Enc)
		if err != nil {
			return fmt.Errorf("failed to decrypt key: %v", err)
		}

		contentKey := &ContentKey{
			KeyID:  keyID,
			Type:   key.GetType().String(),
			Bytes:  decryptedKey,
		}

		if key.GetType() == License_KeyContainer_OPERATOR_SESSION {
			// Handle permissions if this is an operator session key
			// (Implementation depends on your specific needs)
		}

		session.ContentKeys = append(session.ContentKeys, contentKey)
	}

	return nil
}

// decryptKey decrypts a content key using AES-CBC
func decryptKey(encryptedKey, iv, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(iv) != aes.BlockSize {
		return nil, errors.New("invalid IV length")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	decryptedKey := make([]byte, len(encryptedKey))
	mode.CryptBlocks(decryptedKey, encryptedKey)

	// Remove PKCS7 padding
	padding := int(decryptedKey[len(decryptedKey)-1])
	if padding > aes.BlockSize || padding == 0 {
		return nil, errors.New("invalid padding")
	}

	for i := len(decryptedKey) - padding; i < len(decryptedKey); i++ {
		if int(decryptedKey[i]) != padding {
			return nil, errors.New("invalid padding")
		}
	}

	return decryptedKey[:len(decryptedKey)-padding], nil
}

// DeriveKeys derives encryption keys from the session key
func (cdm *CDM) DeriveKeys(message, key []byte) *DerivedKeys {
	encKeyBase := append([]byte("ENCRYPTION\x00"), message...)
	encKeyBase = append(encKeyBase, []byte{0x00, 0x00, 0x00, 0x80}...)

	authKeyBase := append([]byte("AUTHENTICATION\x00"), message...)
	authKeyBase = append(authKeyBase, []byte{0x00, 0x00, 0x02, 0x00}...)

	encKey := append([]byte{0x01}, encKeyBase...)
	authKey1 := append([]byte{0x01}, authKeyBase...)
	authKey2 := append([]byte{0x02}, authKeyBase...)
	authKey3 := append([]byte{0x03}, authKeyBase...)
	authKey4 := append([]byte{0x04}, authKeyBase...)

	encCmacKey := getCMACDigest(encKey, key)
	authCmacKey1 := getCMACDigest(authKey1, key)
	authCmacKey2 := getCMACDigest(authKey2, key)
	authCmacKey3 := getCMACDigest(authKey3, key)
	authCmacKey4 := getCMACDigest(authKey4, key)

	authCombined1 := append(authCmacKey1, authCmacKey2...)
	authCombined2 := append(authCmacKey3, authCmacKey4...)

	return &DerivedKeys{
		Auth1: authCombined1,
		Auth2: authCombined2,
		Enc:   encCmacKey,
	}
}

// GetKeys returns the content keys for a session
func (cdm *CDM) GetKeys(sessionID string) ([]*ContentKey, error) {
	session, exists := cdm.Sessions[sessionID]
	if !exists {
		return nil, errors.New("session not found")
	}
	return session.ContentKeys, nil
}

// getHMACSHA256Digest computes HMAC-SHA256
func getHMACSHA256Digest(data, key []byte) []byte {
	mac := hmac.New(sha256.New, key)
	mac.Write(data)
	return mac.Sum(nil)
}

// getCMACDigest computes CMAC (using AES-CBC as a simplified stand-in)
func getCMACDigest(data, key []byte) []byte {
	// Note: This is a simplified version. Real CMAC is more complex.
	// For a proper implementation, you'd need a dedicated CMAC implementation.
	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	mode := cipher.NewCBCEncrypter(block, make([]byte, aes.BlockSize))
	paddedData := padData(data, aes.BlockSize)
	result := make([]byte, len(paddedData))
	mode.CryptBlocks(result, paddedData)

	// Return the last block as the CMAC
	return result[len(result)-aes.BlockSize:]
}

// padData pads data to the block size
func padData(data []byte, blockSize int) []byte {
	padding := blockSize - (len(data) % blockSize)
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// randInt generates a random integer up to max
func randInt(max int) int {
	n, _ := rand.Int(rand.Reader, big.NewInt(int64(max)))
	return int(n.Int64())
}

// CDMDevice represents a CDM device
type CDMDevice struct {
	DeviceName string
	ClientID   *ClientIdentification
	DeviceKeys *rsa.PrivateKey
	IsAndroid  bool
}

// NewCDMDevice creates a new CDM device
func NewCDMDevice(deviceName string, clientIDBytes, privateKeyBytes, vmpBytes []byte) *CDMDevice {
	// In a real implementation, you would load these from files
	// as shown in the C# code, but we'll simplify for this example
	
	clientID := &ClientIdentification{
		Type: ClientIdentification_KEYBOX.Enum(),
		// Other fields would be initialized here
	}
	
	var privateKey *rsa.PrivateKey
	// Parse private key from bytes if provided
	
	return &CDMDevice{
		DeviceName: deviceName,
		ClientID:   clientID,
		DeviceKeys: privateKey,
		IsAndroid:  true,
	}
}

// Decrypt decrypts data using the device's private key
func (d *CDMDevice) Decrypt(data []byte) ([]byte, error) {
	if d.DeviceKeys == nil {
		return nil, errors.New("no device keys available")
	}
	
	blockSize := d.DeviceKeys.Size()
	var plaintext []byte
	
	for i := 0; i < len(data); i += blockSize {
		end := i + blockSize
		if end > len(data) {
			end = len(data)
		}
		
		chunk, err := rsa.DecryptOAEP(sha1.New(), rand.Reader, d.DeviceKeys, data[i:end], nil)
		if err != nil {
			return nil, err
		}
		
		plaintext = append(plaintext, chunk...)
	}
	
	return plaintext, nil
}

// Sign signs data using the device's private key
func (d *CDMDevice) Sign(data []byte) ([]byte, error) {
	if d.DeviceKeys == nil {
		return nil, errors.New("no device keys available")
	}
	
	hasher := sha1.New()
	hasher.Write(data)
	hash := hasher.Sum(nil)
	
	return rsa.SignPSS(rand.Reader, d.DeviceKeys, crypto.SHA1, hash, nil)
}

// Session represents a DRM session
type Session struct {
	SessionId         []byte
	InitData          interface{} // *WidevineCencHeader or []byte
	Offline           bool
	Device            *CDMDevice
	SessionKey        []byte
	DerivedKeys       *DerivedKeys
	LicenseRequest    []byte
	License           *SignedLicense
	ServiceCertificate *SignedDeviceCertificate
	PrivacyMode       bool
	ContentKeys       []*ContentKey
}

// NewSession creates a new session
func NewSession(sessionID []byte, initData interface{}, device *CDMDevice, offline bool) *Session {
	return &Session{
		SessionId:   sessionID,
		InitData:    initData,
		Offline:     offline,
		Device:      device,
		ContentKeys: make([]*ContentKey, 0),
	}
}

// DerivedKeys contains the derived encryption keys
type DerivedKeys struct {
	Auth1 []byte
	Auth2 []byte
	Enc   []byte
}

// ContentKey represents a decrypted content key
type ContentKey struct {
	KeyID       []byte
	Type        string
	Bytes       []byte
	Permissions []string
}

// PSSHBox represents a PSSH box
type PSSHBox struct {
	KIDs [][]byte
	Data []byte
}

// ParsePSSHBox parses a PSSH box from bytes
func ParsePSSHBox(data []byte) (*PSSHBox, error) {
	if len(data) < 32 {
		return nil, errors.New("PSSH box too short")
	}

	// Skip the first 20 bytes (box size and type)
	stream := bytes.NewReader(data)
	stream.Seek(20, io.SeekStart)

	// Read KID count
	kidCountBytes := make([]byte, 4)
	if _, err := stream.Read(kidCountBytes); err != nil {
		return nil, err
	}

	kidCount := binary.BigEndian.Uint32(kidCountBytes)
	kids := make([][]byte, 0, kidCount)

	for i := 0; i < int(kidCount); i++ {
		kid := make([]byte, 16)
		if _, err := stream.Read(kid); err != nil {
			return nil, err
		}
		kids = append(kids, kid)
	}

	// Read data length
	dataLenBytes := make([]byte, 4)
	if _, err := stream.Read(dataLenBytes); err != nil {
		return nil, err
	}

	dataLen := binary.BigEndian.Uint32(dataLenBytes)
	if dataLen == 0 {
		return &PSSHBox{KIDs: kids}, nil
	}

	// Read data
	psshData := make([]byte, dataLen)
	if _, err := stream.Read(psshData); err != nil {
		return nil, err
	}

	return &PSSHBox{
		KIDs: kids,
		Data: psshData,
	}, nil
}