package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

// ---- parseUsageStats -------------------------------------------------------

type counterPair struct {
	Uplink   int64
	Downlink int64
}

// orderedCounters keeps first-insertion key order like a JS Map.
type orderedCounters struct {
	keys []string
	m    map[string]*counterPair
}

func newOrderedCounters() *orderedCounters {
	return &orderedCounters{m: map[string]*counterPair{}}
}

func (o *orderedCounters) set(key, direction string, value int64) {
	entry, ok := o.m[key]
	if !ok {
		entry = &counterPair{}
		o.m[key] = entry
		o.keys = append(o.keys, key)
	}
	if direction == "uplink" {
		entry.Uplink = value
	} else {
		entry.Downlink = value
	}
}

var (
	inboundStatRe = regexp.MustCompile(`^inbound>>>(.+)>>>traffic>>>(uplink|downlink)$`)
	userStatRe    = regexp.MustCompile(`^user>>>(.+)>>>traffic>>>(uplink|downlink)$`)
)

type parsedStats struct {
	inbounds *orderedCounters
	users    *orderedCounters
}

// parseUsageStats parses the stats helper's JSON output — agent.js
// parseUsageStats(). Zero traffic makes the helper print {"stat":null},
// which is "no data", not a format error.
func parseUsageStats(output string) (*parsedStats, error) {
	invalid := fmt.Errorf("统计接口返回的 JSON stat 格式无效")
	var top map[string]json.RawMessage
	if err := json.Unmarshal([]byte(output), &top); err != nil || top == nil {
		return nil, invalid
	}
	statRaw, present := top["stat"]
	if !present || strings.TrimSpace(string(statRaw)) == "null" {
		return &parsedStats{inbounds: newOrderedCounters(), users: newOrderedCounters()}, nil
	}
	var entries []struct {
		Name  string          `json:"name"`
		Value json.RawMessage `json:"value"`
	}
	if err := json.Unmarshal(statRaw, &entries); err != nil {
		return nil, invalid
	}
	parsed := &parsedStats{inbounds: newOrderedCounters(), users: newOrderedCounters()}
	for _, stat := range entries {
		var bucket *orderedCounters
		var key, direction string
		if m := inboundStatRe.FindStringSubmatch(stat.Name); m != nil {
			bucket, key, direction = parsed.inbounds, m[1], m[2]
		} else if m := userStatRe.FindStringSubmatch(stat.Name); m != nil {
			bucket, key, direction = parsed.users, m[1], m[2]
		} else {
			continue
		}
		value, reason, valid := jsStatValue(stat.Value)
		if !valid {
			return nil, fmt.Errorf("入口计数异常（%s）：%s；请运行 ng-agent doctor 查看原始统计", reason, stat.Name)
		}
		bucket.set(key, direction, value)
	}
	return parsed, nil
}

// ---- queryUsage ------------------------------------------------------------

type usageSample struct {
	Kind     string `json:"kind"`
	ID       string `json:"id"`
	Uplink   int64  `json:"uplink"`
	Downlink int64  `json:"downlink"`
	Epoch    string `json:"epoch"`
}

type usageReport struct {
	Samples      []usageSample `json:"samples"`
	Error        *string       `json:"error"`
	EngineErrors []string      `json:"engineErrors"`
}

func (r *usageReport) toMap() map[string]interface{} {
	return map[string]interface{}{
		"samples":      r.Samples,
		"error":        r.Error,
		"engineErrors": r.EngineErrors,
	}
}

// queryUsage reads sing-box counters through the stats helper binary and
// shapes them into per-node / per-user samples — agent.js queryUsage().
// Callers must hold stateMu.
func queryUsageLocked() (*usageReport, error) {
	all, err := readResources()
	if err != nil {
		return nil, err
	}
	resources := []map[string]interface{}{}
	for _, res := range all {
		meta := resMeta(res)
		if meta != nil && jsTruthy(meta["metricsTag"]) && asString(meta["kind"]) == "node" {
			resources = append(resources, res)
		}
	}
	samples := []usageSample{}
	errorsList := []string{}
	engineErrors := []string{}
	if len(resources) == 0 {
		return &usageReport{Samples: samples, Error: nil, EngineErrors: engineErrors}, nil
	}
	before := serviceEpoch()
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	cmd := exec.CommandContext(ctx, cfg.SingBoxStatsBin, "--server=127.0.0.1:"+cfg.APIPort)
	var outBuf, errBuf bytes.Buffer
	cmd.Stdout = &outBuf
	cmd.Stderr = &errBuf
	runErr := cmd.Run()
	cancel()
	after := serviceEpoch()
	failed := false
	parsed := &parsedStats{inbounds: newOrderedCounters(), users: newOrderedCounters()}
	if before != "" && after != "" && before != after {
		errorsList = append(errorsList, "sing-box restarted during statistics query")
		engineErrors = append(engineErrors, "sing-box")
		failed = true
	} else if runErr != nil {
		detail := errBuf.String()
		if detail == "" && runErr != nil {
			detail = runErr.Error()
		}
		if detail == "" {
			detail = outBuf.String()
		}
		errorsList = append(errorsList, fmt.Sprintf("sing-box statsquery failed: %s",
			truncateRunes(safeError(detail), 160)))
		engineErrors = append(engineErrors, "sing-box")
		failed = true
	}
	if !failed {
		part, err := parseUsageStats(outBuf.String())
		if err != nil {
			errorsList = append(errorsList, "sing-box: "+safeError(err.Error()))
			engineErrors = append(engineErrors, "sing-box")
			failed = true
		} else {
			parsed = part
		}
	}
	if !failed {
		epochValue := before
		if epochValue == "" {
			epochValue = serviceEpoch()
		}
		if epochValue == "" {
			epochValue = "unknown"
		}
		epoch := "sing-box:" + epochValue
		tagNodes := map[string]string{}
		for _, res := range resources {
			meta := resMeta(res)
			tagNodes[asString(meta["metricsTag"])] = asString(meta["nodeId"])
		}
		for _, tag := range parsed.inbounds.keys {
			counters := parsed.inbounds.m[tag]
			nodeID, ok := tagNodes[tag]
			if !ok || nodeID == "" {
				continue
			}
			samples = append(samples, usageSample{Kind: "inbound", ID: nodeID,
				Uplink: counters.Uplink, Downlink: counters.Downlink, Epoch: epoch})
		}
		userRe := regexp.MustCompile(`^ng:(.+)$`)
		for _, name := range parsed.users.keys {
			counters := parsed.users.m[name]
			match := userRe.FindStringSubmatch(name)
			if match == nil {
				continue
			}
			samples = append(samples, usageSample{Kind: "user", ID: match[1],
				Uplink: counters.Uplink, Downlink: counters.Downlink, Epoch: epoch})
		}
		if len(parsed.inbounds.keys) == 0 && len(parsed.users.keys) == 0 {
			errorsList = append(errorsList, "sing-box 入口尚无统计记录")
		}
	}
	var errPtr *string
	if len(errorsList) > 0 {
		joined := strings.Join(errorsList, "; ")
		errPtr = &joined
	}
	return &usageReport{Samples: samples, Error: errPtr, EngineErrors: engineErrors}, nil
}

func queryUsage() (*usageReport, error) {
	stateMu.Lock()
	defer stateMu.Unlock()
	return queryUsageLocked()
}

// ---- flush before config changes ------------------------------------------

// flushUsageBeforeChange persists the running engine's counters before an
// apply/delete restarts sing-box and destroys them — agent.js
// flushUsageBeforeChange(). Callers must hold stateMu.
func flushUsageBeforeChangeLocked() error {
	resources, err := readResources()
	if err != nil {
		return err
	}
	live := false
	for _, res := range resources {
		meta := resMeta(res)
		if meta != nil && jsTruthy(meta["metricsTag"]) && asString(meta["kind"]) == "node" {
			live = true
			break
		}
	}
	if !live || serviceHealth("nexusgate-sing-box").Status != "ready" {
		return nil
	}
	report, err := queryUsageLocked()
	if err != nil {
		return err
	}
	for _, e := range report.EngineErrors {
		if e == "sing-box" {
			return fmt.Errorf("sing-box 重载前无法读取入口计数；请先运行 ng-agent doctor")
		}
	}
	// Await the controller's durable write before the restart.
	_, err = requestJSON("/api/agent/usage", report.toMap())
	return err
}

// flushAllUsageBeforeRestart backs the `flush-usage` CLI action used by
// agent-update.sh.
func flushAllUsageBeforeRestart() error {
	stateMu.Lock()
	defer stateMu.Unlock()
	return flushUsageBeforeChangeLocked()
}

// ---- AnyTLS IP observations ------------------------------------------------

var (
	anytlsCtxRe  = regexp.MustCompile(`\[(\d+)\s+\d+ms\]\s+inbound/anytls\[([^\]]+)\]:\s+(.*)$`)
	anytlsSrcRe  = regexp.MustCompile(`inbound connection from\s+(\[[^\]]+\]|[\d.]+):\d+`)
	anytlsAuthRe = regexp.MustCompile(`\[ng:([^\]]+)\] inbound connection to `)
)

type observation struct {
	CustomerID string `json:"customerId"`
	IP         string `json:"ip"`
}

type observationsResult struct {
	observations []observation
	offset       *int64
}

// readSingBoxObservations tails the sing-box access log from the saved cursor
// and extracts authenticated AnyTLS (customerId, source IP) pairs — agent.js
// readSingBoxObservations(). Callers must hold stateMu.
func readSingBoxObservationsLocked() observationsResult {
	empty := observationsResult{observations: []observation{}, offset: nil}
	info, err := os.Stat(cfg.SingBoxLog)
	if err != nil {
		return empty
	}
	var cursorOffset int64
	if data, err := os.ReadFile(cfg.SingBoxCursor); err == nil {
		var cursor struct {
			Offset int64 `json:"offset"`
		}
		if err := json.Unmarshal(data, &cursor); err == nil {
			cursorOffset = cursor.Offset
		}
	}
	size := info.Size()
	if size < cursorOffset {
		cursorOffset = 0
	}
	start := cursorOffset
	if floor := size - 2*1024*1024; floor > start {
		start = floor
	}
	if start >= size {
		return empty
	}
	f, err := os.Open(cfg.SingBoxLog)
	if err != nil {
		return empty
	}
	defer f.Close()
	buffer := make([]byte, size-start)
	if _, err := f.ReadAt(buffer, start); err != nil {
		return empty
	}
	end := bytes.LastIndex(buffer, []byte("\n"))
	if end < 0 {
		return empty
	}
	resources, err := readResources()
	if err != nil {
		return empty
	}
	entryTags := map[string]bool{}
	for _, res := range resources {
		if meta := resMeta(res); meta != nil && asString(meta["kind"]) == "node" {
			if tag := asString(meta["metricsTag"]); tag != "" {
				entryTags[tag] = true
			}
		}
	}
	sourceByID := map[string]string{}
	uniqueKeys := []string{}
	unique := map[string]observation{}
	for _, line := range strings.Split(string(buffer[:end]), "\n") {
		ctx := anytlsCtxRe.FindStringSubmatch(line)
		if ctx == nil || !entryTags[ctx[2]] {
			continue
		}
		key := ctx[2] + ":" + ctx[1]
		if src := anytlsSrcRe.FindStringSubmatch(ctx[3]); src != nil {
			ip := src[1]
			ip = strings.TrimPrefix(ip, "[")
			ip = strings.TrimSuffix(ip, "]")
			sourceByID[key] = ip
		}
		// Only count authenticated AnyTLS sessions, never unauthenticated probes.
		if authed := anytlsAuthRe.FindStringSubmatch(ctx[3]); authed != nil {
			if ip, ok := sourceByID[key]; ok {
				customerID := authed[1]
				uniqKey := customerID + "|" + ip
				if _, exists := unique[uniqKey]; !exists {
					uniqueKeys = append(uniqueKeys, uniqKey)
				}
				unique[uniqKey] = observation{CustomerID: customerID, IP: ip}
			}
		}
	}
	result := observationsResult{observations: []observation{}}
	for _, k := range uniqueKeys {
		result.observations = append(result.observations, unique[k])
	}
	offset := start + int64(end) + 1
	result.offset = &offset
	return result
}

// ---- periodic loops ----------------------------------------------------------

func usageLoop() {
	for {
		sleepMS(60000)
		func() {
			report, err := queryUsage()
			var body map[string]interface{}
			if err != nil {
				// JS catch fallback: { samples: [], error } (no engineErrors key).
				msg := safeError(err.Error())
				body = map[string]interface{}{"samples": []usageSample{}, "error": msg}
			} else {
				body = report.toMap()
			}
			if _, err := requestJSON("/api/agent/usage", body); err != nil {
				logf("Usage report failed %s", safeError(err.Error()))
			}
		}()
		func() {
			stateMu.Lock()
			defer stateMu.Unlock()
			sing := readSingBoxObservationsLocked()
			if len(sing.observations) > 0 {
				if _, err := requestJSON("/api/agent/observations",
					map[string]interface{}{"observations": sing.observations}); err != nil {
					logf("IP observation report failed %s", safeError(err.Error()))
					return
				}
			}
			// A failed upload must not advance the cursor; the same log
			// lines are retried next cycle (the early return above matches
			// the JS try/catch scope).
			if sing.offset != nil {
				body, _ := marshalJSON(map[string]interface{}{"offset": *sing.offset})
				if err := os.WriteFile(cfg.SingBoxCursor, []byte(body+"\n"), 0o600); err != nil {
					logf("IP observation report failed %s", safeError(err.Error()))
				}
			}
		}()
	}
}