package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/netip"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"mtproxy_checker/internal/allowlist"
	"mtproxy_checker/internal/checker"
	"mtproxy_checker/internal/checkresult"
	"mtproxy_checker/internal/dclist"
	"mtproxy_checker/internal/parseurl"
	"mtproxy_checker/internal/secret"
)

func main() {
	log.SetFlags(log.LstdFlags | log.LUTC)
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

type config struct {
	listFile        string
	checkInterval   time.Duration
	httpAddr        string
	checkTimeout    time.Duration
	dcIDs           []int16
	probe           checker.ProbeMode
	allowedPrefixes []netip.Prefix
	tdlibHelper     string
	tdlibTimeout    time.Duration
}

func loadConfig() (*config, error) {
	listFile := strings.TrimSpace(os.Getenv("MTPROXY_LIST_FILE"))
	if listFile == "" {
		listFile = "/data/proxies.txt"
	}
	intervalStr := strings.TrimSpace(os.Getenv("MTPROXY_CHECK_INTERVAL"))
	if intervalStr == "" {
		intervalStr = "5m"
	}
	interval, err := time.ParseDuration(intervalStr)
	if err != nil || interval <= 0 {
		return nil, fmt.Errorf("MTPROXY_CHECK_INTERVAL: invalid duration %q", intervalStr)
	}
	httpAddr := strings.TrimSpace(os.Getenv("MTPROXY_HTTP_ADDR"))
	if httpAddr == "" {
		httpAddr = ":8080"
	}
	timeoutStr := strings.TrimSpace(os.Getenv("MTPROXY_CHECK_TIMEOUT"))
	if timeoutStr == "" {
		// Fake-TLS + drain + req_pq/resPQ к DC часто >15s на медленных линиях и при многих прокси подряд.
		timeoutStr = "45s"
	}
	checkTimeout, err := time.ParseDuration(timeoutStr)
	if err != nil || checkTimeout <= 0 {
		return nil, fmt.Errorf("MTPROXY_CHECK_TIMEOUT: invalid duration %q", timeoutStr)
	}
	dcIDsRaw := strings.TrimSpace(os.Getenv("MTPROXY_DC_IDS"))
	dcIDs, err := dclist.ParseList(dcIDsRaw)
	if err != nil {
		return nil, fmt.Errorf("MTPROXY_DC_IDS: %w", err)
	}
	if dcIDs == nil {
		dcStr := strings.TrimSpace(os.Getenv("MTPROXY_DC_ID"))
		if dcStr == "" {
			dcStr = "2"
		}
		var dcParsed int64
		_, err = fmt.Sscanf(dcStr, "%d", &dcParsed)
		if err != nil || dcParsed < -32768 || dcParsed > 32767 {
			return nil, fmt.Errorf("MTPROXY_DC_ID: invalid int16 %q", dcStr)
		}
		dcIDs = []int16{int16(dcParsed)}
	}
	allowedRaw := os.Getenv("MTPROXY_ALLOWED_IPS")
	prefixes, err := allowlist.ParseCommaList(allowedRaw)
	if err != nil {
		return nil, err
	}
	probe := checker.ParseProbe(os.Getenv("MTPROXY_PROBE"))
	tdlibHelper := strings.TrimSpace(os.Getenv("MTPROXY_TDLIB_HELPER"))
	tdlibTimeoutStr := strings.TrimSpace(os.Getenv("MTPROXY_TDLIB_TIMEOUT"))
	if tdlibTimeoutStr == "" {
		tdlibTimeoutStr = "45s"
	}
	tdlibTimeout, err := time.ParseDuration(tdlibTimeoutStr)
	if err != nil || tdlibTimeout <= 0 {
		return nil, fmt.Errorf("MTPROXY_TDLIB_TIMEOUT: invalid duration %q", tdlibTimeoutStr)
	}
	return &config{
		listFile:        listFile,
		checkInterval:   interval,
		httpAddr:        httpAddr,
		checkTimeout:    checkTimeout,
		dcIDs:           dcIDs,
		probe:           probe,
		allowedPrefixes: prefixes,
		tdlibHelper:     tdlibHelper,
		tdlibTimeout:    tdlibTimeout,
	}, nil
}

type dcProbeResult struct {
	DC       int16  `json:"dc"`
	OK       bool   `json:"ok"`
	ExitCode int    `json:"exit_code"`
	Error    string `json:"error,omitempty"`
}

type standardProbe struct {
	OK         bool   `json:"ok"`
	ExitCode   int    `json:"exit_code"`
	Error      string `json:"error,omitempty"`
	ParseError string `json:"parse_error,omitempty"`
}

type tdlibProbe struct {
	SkippedReason string `json:"skipped_reason,omitempty"`
	Ran           bool   `json:"ran"`
	OK            bool   `json:"ok"`
	ExitCode      int    `json:"exit_code,omitempty"`
	Error         string `json:"error,omitempty"`
	DurationMs    int64  `json:"duration_ms,omitempty"`
}

type proxyEntry struct {
	RawLine    string          `json:"raw_line"`
	URL        string          `json:"url,omitempty"`
	CheckedAt  string          `json:"checked_at,omitempty"`
	Standard   standardProbe   `json:"standard"`
	TDLib      tdlibProbe      `json:"tdlib"`
	OK         bool            `json:"ok"`
	ExitCode   int             `json:"exit_code"`
	Error      string          `json:"error,omitempty"`
	ParseError string          `json:"parse_error,omitempty"`
	DCs        []dcProbeResult `json:"dcs,omitempty"`
}

type snapshot struct {
	CycleFinishedAt string        `json:"cycle_finished_at"`
	NextCheckAfter  string        `json:"next_check_after,omitempty"`
	Proxies         []proxyEntry  `json:"proxies"`
}

type store struct {
	mu   sync.RWMutex
	data snapshot
}

func (s *store) get() snapshot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := s.data
	out.Proxies = append([]proxyEntry(nil), s.data.Proxies...)
	return out
}

func (s *store) setCycle(entries []proxyEntry, finished time.Time, next time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.data.Proxies = append([]proxyEntry(nil), entries...)
	s.data.CycleFinishedAt = finished.UTC().Format(time.RFC3339)
	if !next.IsZero() {
		s.data.NextCheckAfter = next.UTC().Format(time.RFC3339)
	} else {
		s.data.NextCheckAfter = ""
	}
}

func readProxyLines(path string) ([]string, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var lines []string
	for _, line := range strings.Split(string(b), "\n") {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		lines = append(lines, line)
	}
	return lines, nil
}

func aggregateExitFromDCs(dcs []dcProbeResult) int {
	has3, has4 := false, false
	for _, d := range dcs {
		if d.ExitCode == 3 {
			has3 = true
		}
		if d.ExitCode == 4 {
			has4 = true
		}
	}
	if has3 {
		return 3
	}
	if has4 {
		return 4
	}
	return 1
}

func checkOneLine(ctx context.Context, line string, dcIDs []int16, perDCTimeout time.Duration, probe checker.ProbeMode, tdlibHelper string, tdlibTimeout time.Duration) (ent proxyEntry) {
	now := time.Now().UTC().Format(time.RFC3339)
	ent = proxyEntry{RawLine: line, CheckedAt: now}
	defer func() {
		ent.Standard = standardProbe{OK: ent.OK, ExitCode: ent.ExitCode, Error: ent.Error, ParseError: ent.ParseError}
		ent.TDLib = runTDLibBlock(line, tdlibHelper, tdlibTimeout, probe, ent)
	}()
	if len(dcIDs) == 0 {
		ent.ExitCode = 2
		ent.Error = "no DC ids configured"
		return ent
	}
	t, err := parseurl.ParseTGProxy(line)
	if err != nil {
		ent.ExitCode = 2
		ent.ParseError = err.Error()
		ent.Error = ent.ParseError
		return ent
	}
	ent.URL = line
	parsed, err := secret.Parse(t.Secret)
	if err != nil {
		ent.ExitCode = 2
		ent.ParseError = err.Error()
		ent.Error = ent.ParseError
		return ent
	}
	if len(dcIDs) == 1 {
		dcCtx, cancel := context.WithTimeout(ctx, perDCTimeout)
		err = checker.Check(dcCtx, t.Host, t.Port, parsed, dcIDs[0], &checker.Options{Probe: probe})
		cancel()
		code, msg := checkresult.Classify(err)
		ent.ExitCode = code
		ent.OK = err == nil
		if msg != "" {
			ent.Error = msg
		}
		return ent
	}
	var dcs []dcProbeResult
	anyOK := false
	for _, dc := range dcIDs {
		dcCtx, cancel := context.WithTimeout(ctx, perDCTimeout)
		err := checker.Check(dcCtx, t.Host, t.Port, parsed, dc, &checker.Options{Probe: probe})
		cancel()
		code, msg := checkresult.Classify(err)
		ok := err == nil
		if ok {
			anyOK = true
		}
		dcs = append(dcs, dcProbeResult{DC: dc, OK: ok, ExitCode: code, Error: msg})
	}
	ent.DCs = dcs
	ent.OK = anyOK
	if ent.OK {
		ent.ExitCode = 0
	} else {
		ent.ExitCode = aggregateExitFromDCs(dcs)
		var parts []string
		for _, d := range dcs {
			if d.Error != "" {
				parts = append(parts, fmt.Sprintf("dc%d: %s", d.DC, d.Error))
			} else {
				parts = append(parts, fmt.Sprintf("dc%d: fail", d.DC))
			}
		}
		ent.Error = strings.Join(parts, "; ")
	}
	return ent
}

func runCycle(cfg *config, st *store) {
	lines, err := readProxyLines(cfg.listFile)
	if err != nil {
		finished := time.Now().UTC()
		st.setCycle([]proxyEntry{{
			RawLine:    cfg.listFile,
			ExitCode:   2,
			ParseError: err.Error(),
			Error:      err.Error(),
			CheckedAt:  finished.Format(time.RFC3339),
			Standard:   standardProbe{OK: false, ExitCode: 2, Error: err.Error(), ParseError: err.Error()},
			TDLib:      tdlibProbe{SkippedReason: "cycle aborted (list file error)", Ran: false},
		}}, finished, time.Now().Add(cfg.checkInterval))
		log.Printf("read list file: %v", err)
		return
	}
	entries := make([]proxyEntry, 0, len(lines))
	for _, line := range lines {
		total := cfg.checkTimeout
		if len(cfg.dcIDs) > 1 {
			total = cfg.checkTimeout * time.Duration(len(cfg.dcIDs))
		}
		if cfg.probe == checker.ProbeFast && strings.TrimSpace(cfg.tdlibHelper) != "" {
			total += cfg.tdlibTimeout
		}
		ctx, cancel := context.WithTimeout(context.Background(), total)
		ent := checkOneLine(ctx, line, cfg.dcIDs, cfg.checkTimeout, cfg.probe, cfg.tdlibHelper, cfg.tdlibTimeout)
		cancel()
		entries = append(entries, ent)
	}
	finished := time.Now().UTC()
	next := time.Now().Add(cfg.checkInterval)
	st.setCycle(entries, finished, next)
}

func whitelistMiddleware(prefixes []netip.Prefix, next http.Handler) http.Handler {
	if len(prefixes) == 0 {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			host = r.RemoteAddr
		}
		if len(host) >= 2 && host[0] == '[' && host[len(host)-1] == ']' {
			host = host[1 : len(host)-1]
		}
		addr, err := netip.ParseAddr(host)
		if err != nil || !allowlist.Contains(prefixes, addr) {
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.WriteHeader(http.StatusForbidden)
			_ = json.NewEncoder(w).Encode(map[string]string{"error": "forbidden"})
			return
		}
		next.ServeHTTP(w, r)
	})
}

func run() error {
	cfg, err := loadConfig()
	if err != nil {
		return err
	}
	st := &store{}

	go func() {
		runCycle(cfg, st)
		t := time.NewTicker(cfg.checkInterval)
		defer t.Stop()
		for range t.C {
			runCycle(cfg, st)
		}
	}()

	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		_, _ = w.Write([]byte(`{"status":"ok"}`))
	})
	mux.HandleFunc("/api/v1/proxies", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		snap := st.get()
		enc := json.NewEncoder(w)
		enc.SetIndent("", "  ")
		if err := enc.Encode(snap); err != nil {
			log.Printf("encode json: %v", err)
		}
	})

	handler := whitelistMiddleware(cfg.allowedPrefixes, mux)
	srv := &http.Server{
		Addr:              cfg.httpAddr,
		Handler:           handler,
		ReadHeaderTimeout: 10 * time.Second,
	}

	errCh := make(chan error, 1)
	go func() {
		if cfg.tdlibHelper != "" {
			log.Printf("listening on %s, list=%s interval=%s dcs=%v tdlib_helper=%s tdlib_timeout=%s", cfg.httpAddr, cfg.listFile, cfg.checkInterval, cfg.dcIDs, cfg.tdlibHelper, cfg.tdlibTimeout)
		} else {
			log.Printf("listening on %s, list=%s interval=%s dcs=%v", cfg.httpAddr, cfg.listFile, cfg.checkInterval, cfg.dcIDs)
		}
		errCh <- srv.ListenAndServe()
	}()

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	select {
	case <-sig:
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		return srv.Shutdown(ctx)
	case err := <-errCh:
		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			return err
		}
	}
	return nil
}