// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (c) 2025 Abel Luck <abel@guardianproject.info>

package provider

import (
	"context"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"strings"

	"github.com/hashicorp/terraform-plugin-framework/resource"
	"github.com/hashicorp/terraform-plugin-framework/resource/schema"
	"github.com/hashicorp/terraform-plugin-framework/resource/schema/int64default"
	"github.com/hashicorp/terraform-plugin-framework/resource/schema/planmodifier"
	"github.com/hashicorp/terraform-plugin-framework/resource/schema/stringplanmodifier"
	"github.com/hashicorp/terraform-plugin-framework/types"
	"github.com/hashicorp/terraform-plugin-log/tflog"

	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/lyrebird/common/drbg"
	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/lyrebird/common/ntor"
)

// Ensure provider defined types fully satisfy framework interfaces.
var _ resource.Resource = &TorObfs4StateResource{}

func NewTorObfs4StateResource() resource.Resource {
	return &TorObfs4StateResource{}
}

// TorObfs4StateResource defines the resource implementation.
type TorObfs4StateResource struct{}

// TorObfs4StateResourceModel describes the resource data model.
type TorObfs4StateResourceModel struct {
	Id                        types.String `tfsdk:"id"`
	RsaIdentityPrivateKey     types.String `tfsdk:"rsa_identity_private_key"`
	Ed25519IdentityPrivateKey types.String `tfsdk:"ed25519_identity_private_key"`
	NodeId                    types.String `tfsdk:"node_id"`
	PrivateKey                types.String `tfsdk:"private_key"`
	PublicKey                 types.String `tfsdk:"public_key"`
	DrbgSeed                  types.String `tfsdk:"drbg_seed"`
	IatMode                   types.Int64  `tfsdk:"iat_mode"`
	Certificate               types.String `tfsdk:"certificate"`
	StateJson                 types.String `tfsdk:"state_json"`
	BridgeLine                types.String `tfsdk:"bridge_line"`
}

// obfs4StateJson represents the JSON structure for the state file
type obfs4StateJson struct {
	NodeId     string `json:"node-id"`
	PrivateKey string `json:"private-key"`
	PublicKey  string `json:"public-key"`
	DrbgSeed   string `json:"drbg-seed"`
	IatMode    int    `json:"iat-mode"`
}

func (r *TorObfs4StateResource) Metadata(ctx context.Context, req resource.MetadataRequest, resp *resource.MetadataResponse) {
	resp.TypeName = req.ProviderTypeName + "_obfs4_state"
}

func (r *TorObfs4StateResource) Schema(ctx context.Context, req resource.SchemaRequest, resp *resource.SchemaResponse) {
	resp.Schema = schema.Schema{
		MarkdownDescription: "Generates obfs4 state and certificate for Tor bridges using external relay identity keys",

		Attributes: map[string]schema.Attribute{
			"id": schema.StringAttribute{
				Computed:            true,
				MarkdownDescription: "Resource identifier",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.UseStateForUnknown(),
				},
			},
			"rsa_identity_private_key": schema.StringAttribute{
				Required:            true,
				Sensitive:           true,
				MarkdownDescription: "RSA identity private key in PEM format (from tor_relay_identity_rsa resource)",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.RequiresReplace(),
				},
			},
			"ed25519_identity_private_key": schema.StringAttribute{
				Required:            true,
				Sensitive:           true,
				MarkdownDescription: "Ed25519 identity private key in PEM format (from tor_relay_identity_ed25519 resource)",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.RequiresReplace(),
				},
			},
			"node_id": schema.StringAttribute{
				Computed:            true,
				MarkdownDescription: "20-byte node ID in hex format",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.UseStateForUnknown(),
				},
			},
			"private_key": schema.StringAttribute{
				Computed:            true,
				Sensitive:           true,
				MarkdownDescription: "32-byte Curve25519 private key in hex format",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.UseStateForUnknown(),
				},
			},
			"public_key": schema.StringAttribute{
				Computed:            true,
				MarkdownDescription: "32-byte Curve25519 public key in hex format",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.UseStateForUnknown(),
				},
			},
			"drbg_seed": schema.StringAttribute{
				Computed:            true,
				Sensitive:           true,
				MarkdownDescription: "24-byte DRBG seed in hex format",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.UseStateForUnknown(),
				},
			},
			"iat_mode": schema.Int64Attribute{
				Optional:            true,
				Computed:            true,
				Default:             int64default.StaticInt64(0),
				MarkdownDescription: "Inter-Arrival Time mode (0=none, 1=enabled, 2=paranoid)",
			},
			"certificate": schema.StringAttribute{
				Computed:            true,
				MarkdownDescription: "Base64-encoded certificate for bridge lines",
				PlanModifiers: []planmodifier.String{
					stringplanmodifier.UseStateForUnknown(),
				},
			},
			"state_json": schema.StringAttribute{
				Computed:            true,
				MarkdownDescription: "Complete obfs4 state in JSON format",
			},
			"bridge_line": schema.StringAttribute{
				Computed:            true,
				MarkdownDescription: "Complete bridge line ready for client use (placeholder IP and fingerprint)",
			},
		},
	}
}

func (r *TorObfs4StateResource) Configure(ctx context.Context, req resource.ConfigureRequest, resp *resource.ConfigureResponse) {
	// No configuration needed for this crypto generation resource
}

func (r *TorObfs4StateResource) Create(ctx context.Context, req resource.CreateRequest, resp *resource.CreateResponse) {
	var data TorObfs4StateResourceModel

	// Read Terraform plan data into the model
	resp.Diagnostics.Append(req.Plan.Get(ctx, &data)...)
	if resp.Diagnostics.HasError() {
		return
	}

	// Generate obfs4 state
	err := r.generateObfs4State(ctx, &data)
	if err != nil {
		resp.Diagnostics.AddError("obfs4 State Generation Error",
			fmt.Sprintf("Unable to generate obfs4 state: %s", err))
		return
	}

	// Set ID to a hash of the public key for uniqueness
	data.Id = types.StringValue(fmt.Sprintf("obfs4-%s", data.PublicKey.ValueString()[:16]))

	tflog.Trace(ctx, "created obfs4 state resource")

	// Save data into Terraform state
	resp.Diagnostics.Append(resp.State.Set(ctx, &data)...)
}

func (r *TorObfs4StateResource) Read(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse) {
	var data TorObfs4StateResourceModel

	// Read Terraform prior state data into the model
	resp.Diagnostics.Append(req.State.Get(ctx, &data)...)
	if resp.Diagnostics.HasError() {
		return
	}

	// For this resource, the state is static once created, so no external API calls needed
	// Just ensure the computed fields are still consistent
	if !data.NodeId.IsNull() && !data.PublicKey.IsNull() {
		// Regenerate certificate to ensure consistency
		cert, err := r.generateCertificate(data.NodeId.ValueString(), data.PublicKey.ValueString())
		if err != nil {
			resp.Diagnostics.AddError("Certificate Generation Error",
				fmt.Sprintf("Unable to regenerate certificate: %s", err))
			return
		}
		data.Certificate = types.StringValue(cert)

		// Regenerate state JSON
		stateJson, err := r.generateStateJson(&data)
		if err != nil {
			resp.Diagnostics.AddError("State JSON Generation Error",
				fmt.Sprintf("Unable to regenerate state JSON: %s", err))
			return
		}
		data.StateJson = types.StringValue(stateJson)

		// Regenerate bridge line
		bridgeLine := r.generateBridgeLine(data.Certificate.ValueString(), data.IatMode.ValueInt64())
		data.BridgeLine = types.StringValue(bridgeLine)
	}

	// Save updated data into Terraform state
	resp.Diagnostics.Append(resp.State.Set(ctx, &data)...)
}

func (r *TorObfs4StateResource) Update(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse) {
	var currentState TorObfs4StateResourceModel
	var plannedState TorObfs4StateResourceModel

	// Read current state
	resp.Diagnostics.Append(req.State.Get(ctx, &currentState)...)
	if resp.Diagnostics.HasError() {
		return
	}

	// Read planned state
	resp.Diagnostics.Append(req.Plan.Get(ctx, &plannedState)...)
	if resp.Diagnostics.HasError() {
		return
	}

	// Only iat_mode can be updated; preserve all crypto fields from current state
	currentState.IatMode = plannedState.IatMode

	// Regenerate computed fields with new iat_mode
	stateJson, err := r.generateStateJson(&currentState)
	if err != nil {
		resp.Diagnostics.AddError("State JSON Generation Error",
			fmt.Sprintf("Unable to regenerate state JSON: %s", err))
		return
	}
	currentState.StateJson = types.StringValue(stateJson)

	// Regenerate bridge line with new iat_mode
	bridgeLine := r.generateBridgeLine(currentState.Certificate.ValueString(), currentState.IatMode.ValueInt64())
	currentState.BridgeLine = types.StringValue(bridgeLine)

	tflog.Trace(ctx, "updated obfs4 state resource")

	// Save updated data into Terraform state
	resp.Diagnostics.Append(resp.State.Set(ctx, &currentState)...)
}

func (r *TorObfs4StateResource) Delete(ctx context.Context, req resource.DeleteRequest, resp *resource.DeleteResponse) {
	// No external resources to clean up - this is a stateless resource
	tflog.Trace(ctx, "deleted obfs4 state resource")
}

// generateObfs4State generates all the obfs4 state fields
func (r *TorObfs4StateResource) generateObfs4State(ctx context.Context, data *TorObfs4StateResourceModel) error {
	// Derive node ID from RSA identity key (required)
	nodeId, err := r.deriveNodeIdFromRsaKey(data.RsaIdentityPrivateKey.ValueString())
	if err != nil {
		return fmt.Errorf("failed to derive node ID from RSA key: %w", err)
	}
	data.NodeId = types.StringValue(nodeId.Hex())

	// Generate Curve25519 keypair
	keypair, err := ntor.NewKeypair(false)
	if err != nil {
		return fmt.Errorf("failed to generate keypair: %w", err)
	}
	data.PrivateKey = types.StringValue(keypair.Private().Hex())
	data.PublicKey = types.StringValue(keypair.Public().Hex())

	// Generate DRBG seed
	seed, err := drbg.NewSeed()
	if err != nil {
		return fmt.Errorf("failed to generate DRBG seed: %w", err)
	}
	data.DrbgSeed = types.StringValue(seed.Hex())

	// Set default iat_mode if not specified
	if data.IatMode.IsNull() {
		data.IatMode = types.Int64Value(0)
	}

	// Generate certificate
	cert, err := r.generateCertificate(data.NodeId.ValueString(), data.PublicKey.ValueString())
	if err != nil {
		return fmt.Errorf("failed to generate certificate: %w", err)
	}
	data.Certificate = types.StringValue(cert)

	// Generate state JSON
	stateJson, err := r.generateStateJson(data)
	if err != nil {
		return fmt.Errorf("failed to generate state JSON: %w", err)
	}
	data.StateJson = types.StringValue(stateJson)

	// Generate bridge line
	bridgeLine := r.generateBridgeLine(data.Certificate.ValueString(), data.IatMode.ValueInt64())
	data.BridgeLine = types.StringValue(bridgeLine)

	return nil
}

// generateCertificate creates the base64-encoded certificate from node ID and public key
func (r *TorObfs4StateResource) generateCertificate(nodeIdHex, publicKeyHex string) (string, error) {
	// Decode hex strings to bytes
	nodeIdBytes, err := hex.DecodeString(nodeIdHex)
	if err != nil {
		return "", fmt.Errorf("failed to decode node ID: %w", err)
	}

	publicKeyBytes, err := hex.DecodeString(publicKeyHex)
	if err != nil {
		return "", fmt.Errorf("failed to decode public key: %w", err)
	}

	// Validate lengths
	if len(nodeIdBytes) != 20 {
		return "", fmt.Errorf("node ID must be 20 bytes, got %d", len(nodeIdBytes))
	}
	if len(publicKeyBytes) != 32 {
		return "", fmt.Errorf("public key must be 32 bytes, got %d", len(publicKeyBytes))
	}

	// Concatenate node ID + public key (52 bytes total)
	certBytes := make([]byte, 52)
	copy(certBytes[:20], nodeIdBytes)
	copy(certBytes[20:], publicKeyBytes)

	// Base64 encode and remove padding
	cert := base64.StdEncoding.EncodeToString(certBytes)
	cert = strings.TrimSuffix(cert, "==")

	return cert, nil
}

// generateStateJson creates the JSON representation of the obfs4 state
func (r *TorObfs4StateResource) generateStateJson(data *TorObfs4StateResourceModel) (string, error) {
	state := obfs4StateJson{
		NodeId:     data.NodeId.ValueString(),
		PrivateKey: data.PrivateKey.ValueString(),
		PublicKey:  data.PublicKey.ValueString(),
		DrbgSeed:   data.DrbgSeed.ValueString(),
		IatMode:    int(data.IatMode.ValueInt64()),
	}

	jsonBytes, err := json.Marshal(state)
	if err != nil {
		return "", fmt.Errorf("failed to marshal state JSON: %w", err)
	}

	return string(jsonBytes), nil
}

// generateBridgeLine creates a complete bridge line with placeholder IP and fingerprint
func (r *TorObfs4StateResource) generateBridgeLine(certificate string, iatMode int64) string {
	// Use placeholder values that users can easily replace
	placeholderIP := "<IP:PORT>"
	placeholderFingerprint := "<FINGERPRINT>"

	return fmt.Sprintf("Bridge obfs4 %s %s cert=%s iat-mode=%d",
		placeholderIP, placeholderFingerprint, certificate, iatMode)
}

// deriveNodeIdFromRsaKey derives a 20-byte node ID from an RSA private key PEM
// This follows the Tor specification for relay identity
func (r *TorObfs4StateResource) deriveNodeIdFromRsaKey(rsaPrivateKeyPem string) (*ntor.NodeID, error) {
	// Parse the PEM-encoded RSA private key
	block, _ := pem.Decode([]byte(rsaPrivateKeyPem))
	if block == nil {
		return nil, fmt.Errorf("failed to parse PEM block")
	}

	var privateKey *rsa.PrivateKey
	var err error

	// Try different RSA private key formats
	switch block.Type {
	case "RSA PRIVATE KEY":
		privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	case "PRIVATE KEY":
		parsedKey, err2 := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err2 == nil {
			var ok bool
			privateKey, ok = parsedKey.(*rsa.PrivateKey)
			if !ok {
				return nil, fmt.Errorf("parsed key is not an RSA private key")
			}
		} else {
			err = err2
		}
	default:
		return nil, fmt.Errorf("unsupported PEM block type: %s", block.Type)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to parse RSA private key: %w", err)
	}

	// Extract the public key and encode it
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal public key: %w", err)
	}

	// Generate SHA1 hash of public key (this is the relay fingerprint/node ID)
	hash := sha1.Sum(publicKeyBytes)

	// Create node ID from the first 20 bytes (which is all of SHA1)
	nodeId, err := ntor.NewNodeID(hash[:])
	if err != nil {
		return nil, fmt.Errorf("failed to create node ID: %w", err)
	}

	return nodeId, nil
}