package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"os"
	"path/filepath"

	"neuralhunt/internal/auth"
)

var rawURL = base64.RawURLEncoding

type privateJWK struct {
	Kty    string   `json:"kty"`
	Crv    string   `json:"crv"`
	X      string   `json:"x"`
	Y      string   `json:"y"`
	D      string   `json:"d"`
	Ext    bool     `json:"ext,omitempty"`
	KeyOps []string `json:"key_ops,omitempty"`
	Alg    string   `json:"alg,omitempty"`
	Use    string   `json:"use,omitempty"`
	Kid    string   `json:"kid,omitempty"`
}

type identityFile struct {
	Version    int            `json:"version"`
	PublicJWK  auth.PublicJWK `json:"publicJwk"`
	PrivateJWK privateJWK     `json:"privateJwk"`
}

type encryptedIdentity struct {
	Version    int    `json:"version"`
	Salt       string `json:"salt"`
	IV         string `json:"iv"`
	Ciphertext string `json:"ciphertext"`
}

func pad32Bytes(b []byte) []byte {
	out := make([]byte, 32)
	if len(b) > len(out) {
		b = b[len(b)-len(out):]
	}
	copy(out[len(out)-len(b):], b)
	return out
}

func generateIdentity() (identityFile, *ecdsa.PrivateKey, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return identityFile{}, nil, err
	}
	x := rawURL.EncodeToString(pad32Bytes(key.X.Bytes()))
	y := rawURL.EncodeToString(pad32Bytes(key.Y.Bytes()))
	d := rawURL.EncodeToString(pad32Bytes(key.D.Bytes()))
	pub := auth.PublicJWK{Kty: "EC", Crv: "P-256", X: x, Y: y, Ext: true, KeyOps: []string{"verify"}}
	priv := privateJWK{Kty: "EC", Crv: "P-256", X: x, Y: y, D: d, Ext: true, KeyOps: []string{"sign"}}
	return identityFile{Version: 1, PublicJWK: pub, PrivateJWK: priv}, key, nil
}

func privateKeyFromIdentity(id identityFile) (*ecdsa.PrivateKey, error) {
	if id.Version != 1 || id.PrivateJWK.Kty != "EC" || id.PrivateJWK.Crv != "P-256" || id.PrivateJWK.D == "" {
		return nil, errors.New("unsupported identity; expected version 1 P-256 JWK")
	}
	db, err := rawURL.DecodeString(id.PrivateJWK.D)
	if err != nil {
		return nil, fmt.Errorf("decode private JWK: %w", err)
	}
	d := new(big.Int).SetBytes(db)
	curve := elliptic.P256()
	if d.Sign() <= 0 || d.Cmp(curve.Params().N) >= 0 {
		return nil, errors.New("invalid P-256 private scalar")
	}
	x, y := curve.ScalarBaseMult(pad32Bytes(db))
	if rawURL.EncodeToString(pad32Bytes(x.Bytes())) != id.PublicJWK.X || rawURL.EncodeToString(pad32Bytes(y.Bytes())) != id.PublicJWK.Y {
		return nil, errors.New("identity public/private key mismatch")
	}
	return &ecdsa.PrivateKey{PublicKey: ecdsa.PublicKey{Curve: curve, X: x, Y: y}, D: d}, nil
}

func defaultIdentityPath() string {
	home, err := os.UserHomeDir()
	if err != nil || home == "" {
		return "./neuralhunt-identity.json"
	}
	return filepath.Join(home, ".neuralhunt", "identity.json")
}

func saveIdentity(path string, id identityFile) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	b, err := json.MarshalIndent(id, "", "  ")
	if err != nil {
		return err
	}
	if err := os.WriteFile(path, b, 0o600); err != nil {
		return err
	}
	return os.Chmod(path, 0o600)
}

func loadOrCreateIdentity(path string) (identityFile, *ecdsa.PrivateKey, bool, error) {
	b, err := os.ReadFile(path)
	if err == nil {
		var id identityFile
		if err := json.Unmarshal(b, &id); err != nil {
			return identityFile{}, nil, false, fmt.Errorf("parse identity %q: %w", path, err)
		}
		key, err := privateKeyFromIdentity(id)
		return id, key, false, err
	}
	if !os.IsNotExist(err) {
		return identityFile{}, nil, false, err
	}
	id, key, err := generateIdentity()
	if err != nil {
		return identityFile{}, nil, false, err
	}
	if err := saveIdentity(path, id); err != nil {
		return identityFile{}, nil, false, err
	}
	return id, key, true, nil
}

func readIdentityImport(path, passphrase string) (identityFile, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return identityFile{}, err
	}
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(b, &probe); err != nil {
		return identityFile{}, err
	}
	if _, encrypted := probe["ciphertext"]; !encrypted {
		var id identityFile
		if err := json.Unmarshal(b, &id); err != nil {
			return identityFile{}, err
		}
		if _, err := privateKeyFromIdentity(id); err != nil {
			return identityFile{}, err
		}
		return id, nil
	}
	if passphrase == "" {
		return identityFile{}, errors.New("encrypted browser identity requires --passphrase or NEURALHUNT_IDENTITY_PASSPHRASE")
	}
	var enc encryptedIdentity
	if err := json.Unmarshal(b, &enc); err != nil {
		return identityFile{}, err
	}
	salt, err := rawURL.DecodeString(enc.Salt)
	if err != nil {
		return identityFile{}, err
	}
	iv, err := rawURL.DecodeString(enc.IV)
	if err != nil {
		return identityFile{}, err
	}
	ct, err := rawURL.DecodeString(enc.Ciphertext)
	if err != nil {
		return identityFile{}, err
	}
	key := pbkdf2SHA256([]byte(passphrase), salt, 250000, 32)
	block, err := aes.NewCipher(key)
	if err != nil {
		return identityFile{}, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return identityFile{}, err
	}
	plain, err := gcm.Open(nil, iv, ct, nil)
	if err != nil {
		return identityFile{}, errors.New("identity decrypt failed (wrong passphrase or damaged export)")
	}
	var id identityFile
	if err := json.Unmarshal(plain, &id); err != nil {
		return identityFile{}, err
	}
	if _, err := privateKeyFromIdentity(id); err != nil {
		return identityFile{}, err
	}
	return id, nil
}

func exportBrowserIdentity(path, passphrase string, id identityFile) error {
	if passphrase == "" {
		return errors.New("export requires --passphrase or NEURALHUNT_IDENTITY_PASSPHRASE")
	}
	plain, err := json.Marshal(id)
	if err != nil {
		return err
	}
	salt := make([]byte, 16)
	iv := make([]byte, 12)
	if _, err := rand.Read(salt); err != nil {
		return err
	}
	if _, err := rand.Read(iv); err != nil {
		return err
	}
	key := pbkdf2SHA256([]byte(passphrase), salt, 250000, 32)
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	ct := gcm.Seal(nil, iv, plain, nil)
	enc := encryptedIdentity{Version: 1, Salt: rawURL.EncodeToString(salt), IV: rawURL.EncodeToString(iv), Ciphertext: rawURL.EncodeToString(ct)}
	b, err := json.MarshalIndent(enc, "", "  ")
	if err != nil {
		return err
	}
	if dir := filepath.Dir(path); dir != "." {
		if err := os.MkdirAll(dir, 0o700); err != nil {
			return err
		}
	}
	return os.WriteFile(path, b, 0o600)
}

func pbkdf2SHA256(password, salt []byte, iterations, keyLen int) []byte {
	hLen := sha256.Size
	blocks := (keyLen + hLen - 1) / hLen
	out := make([]byte, 0, blocks*hLen)
	for block := 1; block <= blocks; block++ {
		mac := hmac.New(sha256.New, password)
		mac.Write(salt)
		var n [4]byte
		binary.BigEndian.PutUint32(n[:], uint32(block))
		mac.Write(n[:])
		u := mac.Sum(nil)
		t := append([]byte(nil), u...)
		for i := 1; i < iterations; i++ {
			mac = hmac.New(sha256.New, password)
			mac.Write(u)
			u = mac.Sum(nil)
			for j := range t {
				t[j] ^= u[j]
			}
		}
		out = append(out, t...)
	}
	return out[:keyLen]
}