package main

import (
	"crypto/sha512"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

func main() {
	// Default behavior: process HTML in frontend/build
	target := "frontend/build"
	if len(os.Args) > 1 {
		target = os.Args[1]
	}

	info, err := os.Stat(target)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	if info.IsDir() {
		fmt.Printf("Generating SRI hashes for assets in %s...\n", target)
		err = filepath.Walk(target, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() && strings.HasSuffix(path, ".html") {
				return processHTML(path, target)
			}
			return nil
		})
	} else {
		// If a single file is provided (like verifier.ts), process it specially
		fmt.Printf("Updating SRI hashes in %s...\n", target)
		err = processSourceFile(target)
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

func processHTML(htmlPath, buildDir string) error {
	content, err := os.ReadFile(filepath.Clean(htmlPath))
	if err != nil {
		return err
	}

	updated := string(content)
	scriptRegex := regexp.MustCompile(`<(script|link)\s+([^>]*)(src|href)=["'](/[^"']+)["']([^>]*)>`)

	matches := scriptRegex.FindAllStringSubmatch(updated, -1)
	for _, match := range matches {
		tag := match[0]
		attr := match[3]
		url := match[4]

		if !strings.HasPrefix(url, "/") || strings.HasPrefix(url, "//") {
			continue
		}

		if strings.Contains(tag, "integrity=") {
			continue
		}

		filePath := filepath.Join(buildDir, url)
		hash, err := calculateSHA384(filePath)
		if err != nil {
			fmt.Printf("  Skipping %s: %v\n", url, err)
			continue
		}

		integrity := fmt.Sprintf("integrity=\"sha384-%s\" crossorigin=\"anonymous\"", hash)
		newTag := strings.Replace(tag, fmt.Sprintf("%s=\"%s\"", attr, url), fmt.Sprintf("%s=\"%s\" %s", attr, url, integrity), 1)
		updated = strings.Replace(updated, tag, newTag, 1)
		fmt.Printf("  Added SRI to %s (%s)\n", url, htmlPath)
	}

	if updated != string(content) {
		return os.WriteFile(filepath.Clean(htmlPath), []byte(updated), 0644)
	}
	return nil
}

func processSourceFile(sourcePath string) error {
	content, err := os.ReadFile(filepath.Clean(sourcePath))
	if err != nil {
		return err
	}

	updated := string(content)

	// We need a way to map the asset to the SRI. For verifier.ts, we know the assets.
	assets := map[string]string{
		"wasm_exec.js":  "frontend/static/verifier/wasm_exec.js",
		"verifier.wasm": "frontend/static/verifier/verifier.wasm",
	}

	for assetName, assetPath := range assets {
		hash, err := calculateSHA384(assetPath)
		if err != nil {
			fmt.Printf("  Warning: could not calculate hash for %s: %v\n", assetName, err)
			continue
		}

		// Find the line that mentions the asset and update the NEXT integrity string
		assetEscaped := strings.ReplaceAll(assetName, ".", "\\.")
		assetPattern := regexp.MustCompile(`['"](/verifier/)?` + assetEscaped + `['"][\s\S]*?sha384-([^'"]+)`)

		matches := assetPattern.FindAllStringSubmatchIndex(updated, -1)
		// Process from end to start to not mess up indices
		for i := len(matches) - 1; i >= 0; i-- {
			match := matches[i]
			oldHashStart, oldHashEnd := match[4], match[5]
			oldHash := updated[oldHashStart:oldHashEnd]

			if oldHash != hash {
				updated = updated[:oldHashStart] + hash + updated[oldHashEnd:]
				fmt.Printf("  Updated SRI for %s in %s\n", assetName, sourcePath)
			}
		}
	}

	if updated != string(content) {
		return os.WriteFile(filepath.Clean(sourcePath), []byte(updated), 0644)
	}
	return nil
}

func calculateSHA384(path string) (string, error) {
	f, err := os.Open(filepath.Clean(path))
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := sha512.New384()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(h.Sum(nil)), nil
}