package main

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

// PackageMeta - alias для удобства
type PackageMeta map[string]interface{}

var (
	targetURLStr = "https://registry.npmjs.org"
	cutoffDate   time.Time
)

func main() {
	dateFlag := flag.String("date", "", "Cutoff date in YYYY-MM-DD format (e.g., 2023-12-01)")
	portFlag := flag.String("port", "8080", "Proxy server port")
	flag.Parse()

	if *dateFlag == "" {
		fmt.Println("Usage: nptm --date=2023-12-01")
		os.Exit(1)
	}

	var err error
	cutoffDate, err = time.Parse("2006-01-02", *dateFlag)
	if err != nil {
		log.Fatalf("Invalid date format: %v", err)
	}
	cutoffDate = cutoffDate.Add(23*time.Hour + 59*time.Minute + 59*time.Second)

	log.Printf("Starting npm proxy on :%s", *portFlag)
	log.Printf("Time Machine Active! Packages after %s will be hidden.", cutoffDate.Format(time.RFC3339))
	
	targetURL, _ := url.Parse(targetURLStr)
	proxy := httputil.NewSingleHostReverseProxy(targetURL)

	originalDirector := proxy.Director
	proxy.Director = func(req *http.Request) {
		originalDirector(req)
		req.Host = targetURL.Host
		req.Header.Set("Accept-Encoding", "identity")
		req.Header.Set("Accept", "application/json") 
	}

	proxy.ModifyResponse = func(resp *http.Response) error {
		if resp.StatusCode != http.StatusOK {
			return nil
		}

		contentType := resp.Header.Get("Content-Type")
		if !strings.Contains(contentType, "json") {
			return nil
		}

		var bodyReader io.ReadCloser = resp.Body
		var isGzipped bool

		if resp.Header.Get("Content-Encoding") == "gzip" {
			gzReader, err := gzip.NewReader(resp.Body)
			if err != nil {
				// Если не смогли распаковать, отдаем как есть
				// (скорее всего будет ошибка у клиента, но мы сделали что могли)
				resp.Body.Close()
				return nil 
			}
			bodyReader = gzReader
			isGzipped = true
		}

		bodyBytes, err := io.ReadAll(bodyReader)
		bodyReader.Close()
		if err != nil {
			return err
		}

		// Вспомогательная функция для возврата тела
		// обратно в ответ с коррекцией заголовков
		setBody := func(data []byte) {
			resp.Body = io.NopCloser(bytes.NewReader(data))
			resp.ContentLength = int64(len(data))
			resp.Header.Set("Content-Length", strconv.Itoa(len(data)))
			if isGzipped {
				resp.Header.Del("Content-Encoding")
			}
			// Удаляем ETag, так как мы (возможно) трогали тело
			// (даже просто распаковка меняет байты)
			resp.Header.Del("ETag")
		}

		var meta PackageMeta
		if err := json.Unmarshal(bodyBytes, &meta); err != nil {
			// Если это не JSON, но мы его распаковали из Gzip, 
			// нужно вернуть распакованные данные с правильным Content-Length
			setBody(bodyBytes)
			return nil
		}

		times, okT := meta["time"].(map[string]interface{})
		versions, okV := meta["versions"].(map[string]interface{})

		if !okT || !okV {
			// Аналогично, возвращаем распакованное тело с правильной длиной
			setBody(bodyBytes)
			return nil
		}

		modifiedMeta, err := filterPackageVersions(meta, times, versions)
		if err != nil {
			log.Printf("Filter error: %v", err)
			setBody(bodyBytes)
			return nil
		}

		finalVersions, _ := modifiedMeta["versions"].(map[string]interface{})
		if len(finalVersions) == 0 {
			resp.StatusCode = http.StatusNotFound
			notFoundBody := []byte(fmt.Sprintf(`{"error": "Not found. All versions published after %s"}`, cutoffDate.Format("2006-01-02")))
			
			// Сбрасываем флаг gzip, так как это наше новое тело
			isGzipped = false 
			resp.Header.Del("Content-Encoding")
			
			resp.Body = io.NopCloser(bytes.NewReader(notFoundBody))
			resp.ContentLength = int64(len(notFoundBody))
			resp.Header.Set("Content-Length", strconv.Itoa(len(notFoundBody)))
			return nil
		}

		newBody, err := json.Marshal(modifiedMeta)
		if err != nil {
			return err
		}

		// Устанавливаем итоговое отфильтрованное тело
		setBody(newBody)
		resp.Header.Set("X-NPM-Time-Machine", "true")

		return nil
	}

	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
		// Игнорируем ошибки "context canceled" (клиент отключился)
		if err != nil && !strings.Contains(err.Error(), "context canceled") {
			log.Printf("Proxy error for %s: %v", r.URL.Path, err)
		}
		w.WriteHeader(http.StatusBadGateway)
	}

	log.Fatal(http.ListenAndServe(":"+*portFlag, proxy))
}

func filterPackageVersions(meta PackageMeta, times map[string]interface{}, versions map[string]interface{}) (PackageMeta, error) {
	toDelete := []string{}
	var newLatestVer string
	var newLatestTime time.Time

	for ver, timeVal := range times {
		if ver == "created" || ver == "modified" {
			continue
		}

		timeStr, ok := timeVal.(string)
		if !ok {
			continue
		}

		pubDate, err := time.Parse(time.RFC3339, timeStr)
		if err != nil {
			continue
		}

		if pubDate.After(cutoffDate) {
			toDelete = append(toDelete, ver)
		} else {
			// Ищем самую свежую по дате среди оставшихся
			if pubDate.After(newLatestTime) {
				newLatestTime = pubDate
				newLatestVer = ver
			}
		}
	}

	for _, ver := range toDelete {
		delete(versions, ver)
		delete(times, ver)
	}

	if distTags, ok := meta["dist-tags"].(map[string]interface{}); ok {
		// Если нашли кандидата на latest
		if newLatestVer != "" {
			currentLatest, _ := distTags["latest"].(string)
			// Если текущий latest удален или (важно!) если по времени он в будущем
			// (но не удален по какой-то причине), заменяем его. Самая простая проверка
			// есть ли currentLatest в versions.
			if _, exists := versions[currentLatest]; !exists {
				distTags["latest"] = newLatestVer
			} else {
				// Дополнительная проверка: иногда latest не самый последний по времени,
				// но в рамках TimeMachine мы, скорее всего, хотим видеть последний
				// доступный на ту дату. Однако, безопаснее менять latest только если
				// старый удален. 
			}
		}

		for tag, verInterface := range distTags {
			verStr, ok := verInterface.(string)
			if !ok {
				continue
			}
			if _, exists := versions[verStr]; !exists {
				delete(distTags, tag)
			}
		}
	}

	if !newLatestTime.IsZero() {
		times["modified"] = newLatestTime.Format(time.RFC3339)
	}

	return meta, nil
}