lingma-proxy-compose/internal/bootstrap/lingma.go

package bootstrap

import (
	"archive/zip"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

type Config struct {
	Enabled              bool
	SourceType           string
	VSIXURL              string
	MarketplacePublisher string
	MarketplaceExtension string
	OutputDir            string
	BinaryPath           string
	AlwaysRefresh        bool
	ForceRefresh         bool
	HTTPTimeout          time.Duration
}

type Result struct {
	Enabled        bool   `json:"enabled"`
	Source         string `json:"source,omitempty"`
	Version        string `json:"version,omitempty"`
	URL            string `json:"url,omitempty"`
	BinaryPath     string `json:"binary_path,omitempty"`
	ReleaseDir     string `json:"release_dir,omitempty"`
	MarkerPath     string `json:"marker_path,omitempty"`
	Downloaded     bool   `json:"downloaded"`
	ReusedExisting bool   `json:"reused_existing"`
	Message        string `json:"message,omitempty"`
}

type marker struct {
	Source       string `json:"source"`
	URL          string `json:"url"`
	Version      string `json:"version"`
	DownloadedAt int64  `json:"downloaded_at"`
	NestedZip    string `json:"nested_zip"`
	Member       string `json:"member"`
	ReleaseRoot  string `json:"release_root"`
	Size         int    `json:"size"`
	Publisher    string `json:"publisher,omitempty"`
	Extension    string `json:"extension,omitempty"`
}

func Ensure(cfg Config) (Result, error) {
	result := Result{Enabled: cfg.Enabled}
	if !cfg.Enabled {
		result.Message = "bootstrap disabled"
		return result, nil
	}
	cfg = normalize(cfg)
	if cfg.BinaryPath == "" {
		return result, fmt.Errorf("bootstrap binary path is required")
	}
	if err := os.MkdirAll(filepath.Dir(cfg.BinaryPath), 0o755); err != nil {
		return result, err
	}
	if err := os.MkdirAll(cfg.OutputDir, 0o755); err != nil {
		return result, err
	}

	resolvedURL := strings.TrimSpace(cfg.VSIXURL)
	version := ""
	if cfg.SourceType == "marketplace" {
		url, ver, err := queryMarketplaceLatestVSIX(cfg)
		if err == nil {
			resolvedURL = url
			version = ver
		} else if resolvedURL == "" {
			return result, err
		}
	}
	if resolvedURL == "" {
		return result, fmt.Errorf("bootstrap VSIX URL is empty")
	}

	result.Source = cfg.SourceType
	result.URL = resolvedURL
	result.Version = version
	result.BinaryPath = cfg.BinaryPath
	result.MarkerPath = filepath.Join(filepath.Dir(cfg.BinaryPath), ".lingma-bootstrap.json")

	oldMarker := readMarker(result.MarkerPath)
	releaseDir := releaseDirForOutput(cfg.OutputDir, oldMarker.ReleaseRoot)
	result.ReleaseDir = releaseDir
	ready := hasRequiredAssets(releaseDir)

	if fileExists(cfg.BinaryPath) && ready && !cfg.ForceRefresh {
		if !cfg.AlwaysRefresh || (version != "" && oldMarker.Version == version) {
			if err := os.Chmod(cfg.BinaryPath, 0o755); err != nil {
				return result, err
			}
			result.ReusedExisting = true
			result.Message = "reused existing Lingma binary"
			return result, nil
		}
	}

	vsixBytes, err := downloadVSIX(resolvedURL, cfg.HTTPTimeout)
	if err != nil {
		if fileExists(cfg.BinaryPath) {
			result.ReusedExisting = true
			result.Message = fmt.Sprintf("download failed, reused existing Lingma: %v", err)
			return result, os.Chmod(cfg.BinaryPath, 0o755)
		}
		return result, err
	}
	result.Downloaded = true

	nestedName, nestedBytes, err := extractNestedZip(vsixBytes)
	if err != nil {
		return result, err
	}
	memberName, binaryBytes, releaseRoot, releaseFiles, err := extractRelease(nestedBytes)
	if err != nil {
		return result, err
	}
	releaseDir = releaseDirForOutput(cfg.OutputDir, releaseRoot)
	result.ReleaseDir = releaseDir
	if err := os.RemoveAll(releaseDir); err != nil {
		return result, err
	}
	if err := writeReleaseFiles(releaseDir, releaseRoot, releaseFiles); err != nil {
		return result, err
	}
	if err := os.WriteFile(cfg.BinaryPath, binaryBytes, 0o755); err != nil {
		return result, err
	}
	if !hasRequiredAssets(releaseDir) {
		return result, fmt.Errorf("extension assets missing after extraction under %s", releaseDir)
	}

	mk := marker{
		Source:       cfg.SourceType,
		URL:          resolvedURL,
		Version:      version,
		DownloadedAt: time.Now().Unix(),
		NestedZip:    nestedName,
		Member:       memberName,
		ReleaseRoot:  releaseRoot,
		Size:         len(binaryBytes),
		Publisher:    cfg.MarketplacePublisher,
		Extension:    cfg.MarketplaceExtension,
	}
	if err := writeMarker(result.MarkerPath, mk); err != nil {
		return result, err
	}
	result.Message = fmt.Sprintf("downloaded Lingma %s", versionOrUnknown(version))
	return result, nil
}

type marketplaceResponse struct {
	Results []struct {
		Extensions []struct {
			Versions []struct {
				Version string `json:"version"`
				Files   []struct {
					AssetType string `json:"assetType"`
					Source    string `json:"source"`
				} `json:"files"`
			} `json:"versions"`
		} `json:"extensions"`
	} `json:"results"`
}

type zipEntry struct {
	Name string
	Data []byte
	Mode os.FileMode
}

func normalize(cfg Config) Config {
	cfg.SourceType = strings.ToLower(strings.TrimSpace(cfg.SourceType))
	if cfg.SourceType == "" {
		cfg.SourceType = "marketplace"
	}
	if cfg.OutputDir == "" {
		cfg.OutputDir = filepath.Join(filepath.Dir(cfg.BinaryPath), "release")
	}
	if cfg.HTTPTimeout <= 0 {
		cfg.HTTPTimeout = 30 * time.Second
	}
	if strings.TrimSpace(cfg.MarketplacePublisher) == "" {
		cfg.MarketplacePublisher = "Alibaba-Cloud"
	}
	if strings.TrimSpace(cfg.MarketplaceExtension) == "" {
		cfg.MarketplaceExtension = "tongyi-lingma"
	}
	if strings.TrimSpace(cfg.VSIXURL) == "" {
		cfg.VSIXURL = "https://tongyi-code.oss-cn-hangzhou.aliyuncs.com/vscode/tongyi-lingma-latest.vsix"
	}
	return cfg
}

func queryMarketplaceLatestVSIX(cfg Config) (string, string, error) {
	payload := map[string]any{
		"filters": []any{map[string]any{
			"criteria": []any{
				map[string]any{"filterType": 7, "value": cfg.MarketplacePublisher + "." + cfg.MarketplaceExtension},
				map[string]any{"filterType": 8, "value": "Microsoft.VisualStudio.Code"},
			},
			"pageNumber": 1,
			"pageSize":   1,
			"sortBy":     0,
			"sortOrder":  0,
		}},
		"assetTypes": []any{},
		"flags":      950,
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return "", "", err
	}
	req, err := http.NewRequest(http.MethodPost, "https://marketplace.visualstudio.com/_apis/public/gallery/extensionquery", bytes.NewReader(body))
	if err != nil {
		return "", "", err
	}
	req.Header.Set("accept", "application/json;api-version=3.0-preview.1")
	req.Header.Set("content-type", "application/json")
	req.Header.Set("x-market-client-id", "VSCode 1.115.0")
	client := &http.Client{Timeout: cfg.HTTPTimeout}
	resp, err := client.Do(req)
	if err != nil {
		return "", "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return "", "", fmt.Errorf("marketplace query status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	var parsed marketplaceResponse
	if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
		return "", "", err
	}
	if len(parsed.Results) == 0 || len(parsed.Results[0].Extensions) == 0 || len(parsed.Results[0].Extensions[0].Versions) == 0 {
		return "", "", fmt.Errorf("no extension found from marketplace")
	}
	version := parsed.Results[0].Extensions[0].Versions[0].Version
	for _, file := range parsed.Results[0].Extensions[0].Versions[0].Files {
		if file.AssetType == "Microsoft.VisualStudio.Services.VSIXPackage" && strings.TrimSpace(file.Source) != "" {
			return file.Source, version, nil
		}
	}
	if version == "" {
		return "", "", fmt.Errorf("no version/vsix url found from marketplace")
	}
	fallback := fmt.Sprintf("https://marketplace.visualstudio.com/_apis/public/gallery/publishers/%s/vsextensions/%s/%s/vspackage", cfg.MarketplacePublisher, cfg.MarketplaceExtension, version)
	return fallback, version, nil
}

func downloadVSIX(url string, timeout time.Duration) ([]byte, error) {
	client := &http.Client{Timeout: timeout}
	resp, err := client.Get(url)
	if err != nil {
		return nil, fmt.Errorf("download VSIX: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return nil, fmt.Errorf("download VSIX status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return data, nil
}

func extractNestedZip(vsix []byte) (string, []byte, error) {
	reader, err := zip.NewReader(bytes.NewReader(vsix), int64(len(vsix)))
	if err != nil {
		return "", nil, err
	}
	candidates := make([]*zip.File, 0)
	for _, file := range reader.File {
		name := file.Name
		if strings.HasPrefix(name, "extension/dist/bin/") && strings.HasSuffix(name, ".zip") && strings.Contains(name, "lingma-") {
			candidates = append(candidates, file)
		}
	}
	if len(candidates) == 0 {
		return "", nil, fmt.Errorf("no lingma-*.zip found in VSIX")
	}
	sort.Slice(candidates, func(i, j int) bool { return candidates[i].Name < candidates[j].Name })
	picked := candidates[len(candidates)-1]
	rc, err := picked.Open()
	if err != nil {
		return "", nil, err
	}
	defer rc.Close()
	data, err := io.ReadAll(rc)
	if err != nil {
		return "", nil, err
	}
	return picked.Name, data, nil
}

func extractRelease(inner []byte) (string, []byte, string, []zipEntry, error) {
	reader, err := zip.NewReader(bytes.NewReader(inner), int64(len(inner)))
	if err != nil {
		return "", nil, "", nil, err
	}
	binaryPath := pickLingmaBinaryPath(reader.File)
	if binaryPath == "" {
		return "", nil, "", nil, fmt.Errorf("Lingma binary not found inside nested zip")
	}
	releaseRoot := inferReleaseRoot(binaryPath)
	entries := make([]zipEntry, 0, len(reader.File))
	var binaryBytes []byte
	for _, file := range reader.File {
		if file.FileInfo().IsDir() {
			continue
		}
		if releaseRoot != "" && !strings.HasPrefix(file.Name, releaseRoot+"/") {
			continue
		}
		rc, err := file.Open()
		if err != nil {
			return "", nil, "", nil, err
		}
		data, err := io.ReadAll(rc)
		rc.Close()
		if err != nil {
			return "", nil, "", nil, err
		}
		entries = append(entries, zipEntry{Name: file.Name, Data: data, Mode: file.Mode()})
		if file.Name == binaryPath {
			binaryBytes = data
		}
	}
	if len(binaryBytes) == 0 {
		return "", nil, "", nil, fmt.Errorf("Lingma binary bytes missing from nested zip")
	}
	return binaryPath, binaryBytes, releaseRoot, entries, nil
}

func pickLingmaBinaryPath(files []*zip.File) string {
	preferredSuffix := platformBinarySuffix()
	for _, file := range files {
		if strings.HasSuffix(file.Name, preferredSuffix) {
			return file.Name
		}
	}
	for _, file := range files {
		if strings.HasSuffix(file.Name, "/Lingma") || file.Name == "Lingma" {
			return file.Name
		}
	}
	return ""
}

func platformBinarySuffix() string {
	if runtime.GOOS == "linux" && runtime.GOARCH == "amd64" {
		return "x86_64_linux/Lingma"
	}
	if runtime.GOOS == "linux" && runtime.GOARCH == "arm64" {
		return "arm64_linux/Lingma"
	}
	if runtime.GOOS == "darwin" && runtime.GOARCH == "arm64" {
		return "arm64_darwin/Lingma"
	}
	if runtime.GOOS == "darwin" {
		return "x86_64_darwin/Lingma"
	}
	if runtime.GOOS == "windows" && runtime.GOARCH == "arm64" {
		return "arm64_windows/Lingma"
	}
	if runtime.GOOS == "windows" {
		return "x86_64_windows/Lingma"
	}
	return "/Lingma"
}

func inferReleaseRoot(memberPath string) string {
	parts := strings.Split(memberPath, "/")
	platforms := []string{"x86_64_linux", "arm64_linux", "x86_64_darwin", "arm64_darwin", "x86_64_windows", "arm64_windows"}
	for _, platform := range platforms {
		for i, part := range parts {
			if part == platform && i > 0 {
				return strings.Join(parts[:i], "/")
			}
		}
	}
	if len(parts) > 1 {
		return parts[0]
	}
	return ""
}

func writeReleaseFiles(releaseDir string, releaseRoot string, files []zipEntry) error {
	prefix := ""
	if releaseRoot != "" {
		prefix = releaseRoot + "/"
	}
	for _, entry := range files {
		rel := entry.Name
		if prefix != "" {
			rel = strings.TrimPrefix(rel, prefix)
		}
		if rel == "" {
			continue
		}
		dest := filepath.Join(releaseDir, filepath.FromSlash(rel))
		if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
			return err
		}
		mode := entry.Mode
		if mode == 0 {
			mode = 0o644
		}
		if strings.HasSuffix(rel, "/Lingma") || filepath.Base(rel) == "Lingma" {
			mode = 0o755
		}
		if err := os.WriteFile(dest, entry.Data, mode); err != nil {
			return err
		}
	}
	return nil
}

func releaseDirForOutput(outputDir string, releaseRoot string) string {
	name := strings.TrimSpace(releaseRoot)
	if name == "" {
		return outputDir
	}
	return filepath.Join(outputDir, filepath.FromSlash(name))
}

func hasRequiredAssets(releaseDir string) bool {
	info, err := os.Stat(filepath.Join(releaseDir, "extension", "main.js"))
	return err == nil && !info.IsDir()
}

func readMarker(path string) marker {
	body, err := os.ReadFile(path)
	if err != nil {
		return marker{}
	}
	var mk marker
	if err := json.Unmarshal(body, &mk); err != nil {
		return marker{}
	}
	return mk
}

func writeMarker(path string, mk marker) error {
	body, err := json.MarshalIndent(mk, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, body, 0o644)
}

func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && !info.IsDir()
}

func versionOrUnknown(version string) string {
	if strings.TrimSpace(version) == "" {
		return "unknown version"
	}
	return version
}