package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"
	"time"

	"latprobe/internal/probe"
)

const usageText = `latprobe — measure per-phase HTTP request latency

Usage:
  latprobe [flags] <url> [url ...]

Flags:
  -n, --count int        Number of requests per URL (default 1)
      --timeout duration Request timeout, e.g. 10s, 500ms (default 10s)
      --fail             Exit non-zero on HTTP status >= 400 (exit code 6)
      --json             Output results as JSON instead of text
  -h, --help             Show this help

Exit codes:
  0  All probes succeeded
  1  Usage error
  2  DNS resolution failure
  3  Connection failure
  4  Timeout
  5  TLS handshake failure
  6  HTTP status >= 400 (only with --fail)

Examples:
  latprobe https://example.com
  latprobe -n 5 https://example.com https://www.google.com
  latprobe --timeout 2s https://slow-host.example.com
  latprobe --fail https://example.com
  latprobe --json https://example.com | jq .
`

// exit codes
const (
	exitOK      = 0
	exitUsage   = 1
	exitDNS     = 2
	exitConnect = 3
	exitTimeout = 4
	exitTLS     = 5
	exitHTTP    = 6
)

func failPhaseCode(fp string) int {
	switch fp {
	case "dns":
		return exitDNS
	case "timeout":
		return exitTimeout
	case "tls":
		return exitTLS
	default:
		return exitConnect
	}
}

func main() {
	count := flag.Int("count", 1, "number of requests per URL")
	flag.IntVar(count, "n", 1, "number of requests per URL (shorthand)")
	timeout := flag.Duration("timeout", 10*time.Second, "request timeout per sample")
	fail := flag.Bool("fail", false, "exit non-zero on HTTP status >= 400")
	jsonOut := flag.Bool("json", false, "output results as JSON instead of text")

	flag.Usage = func() { fmt.Fprint(os.Stderr, usageText) }
	flag.Parse()

	urls := flag.Args()
	if len(urls) == 0 {
		fmt.Fprint(os.Stderr, usageText)
		os.Exit(exitUsage)
	}

	opts := probe.Options{Timeout: *timeout}

	worstCode := exitOK
	var jsonEntries []jsonEntry

	for i, rawURL := range urls {
		succeeded, failures := runSamples(rawURL, *count, opts)

		// determine exit code contribution from this URL
		for _, f := range failures {
			if c := failPhaseCode(f.phase); c > worstCode {
				worstCode = c
			}
		}
		if *fail {
			for _, r := range succeeded {
				if r.StatusCode >= 400 && exitHTTP > worstCode {
					worstCode = exitHTTP
				}
			}
		}

		if *jsonOut {
			jsonEntries = append(jsonEntries, buildJSONEntry(rawURL, succeeded, failures))
			continue
		}

		if i > 0 {
			fmt.Println()
		}
		printURL(rawURL, succeeded, failures, *count, *fail)
	}

	if *jsonOut {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		if err := enc.Encode(jsonEntries); err != nil {
			fmt.Fprintf(os.Stderr, "json encode: %v\n", err)
			os.Exit(exitConnect)
		}
	}

	os.Exit(worstCode)
}

// ── sampling ──────────────────────────────────────────────────────────────────

type failItem struct {
	phase   string
	message string
}

func runSamples(rawURL string, count int, opts probe.Options) (succeeded []probe.Result, failures []failItem) {
	for range count {
		r := probe.Measure(rawURL, opts)
		if r.Err != nil {
			failures = append(failures, failItem{
				phase:   r.FailPhase,
				message: unwrapMsg(r.Err),
			})
		} else {
			succeeded = append(succeeded, r)
		}
	}
	return
}

func unwrapMsg(err error) string {
	var urlErr *url.Error
	if errors.As(err, &urlErr) {
		return urlErr.Err.Error()
	}
	return err.Error()
}

// ── text output ───────────────────────────────────────────────────────────────

func printURL(rawURL string, succeeded []probe.Result, failures []failItem, total int, fail bool) {
	nOK := len(succeeded)
	nFail := len(failures)

	switch {
	case nFail == 0 && total == 1:
		// single sample, full success
		printResult(succeeded[0], fail)

	case nFail == 0:
		// multi-sample, all succeeded
		printAggregate(probe.Summarize(succeeded), nil, fail)

	case nOK == 0:
		// all failed — show header + partial phases from last failure result
		header := fmt.Sprintf("%s  (FAILED", rawURL)
		if total > 1 {
			header += fmt.Sprintf(", 0/%d succeeded", total)
		}
		header += ")"
		fmt.Println(header)
		// re-run just to get partial phases from the last failure
		last := probe.Measure(rawURL, probe.Options{Timeout: 1 * time.Millisecond})
		// use the first failure's phase data instead (stored in failures[0])
		// we can't recover partial timing here, so skip phases and go straight to errors
		_ = last
		printFailureSummary(failures)

	default:
		// mixed: some succeeded, some failed
		printAggregate(probe.Summarize(succeeded), failures, fail)
	}
}

func printResult(r probe.Result, fail bool) {
	status := fmt.Sprintf("%d", r.StatusCode)
	if fail && r.StatusCode >= 400 {
		status += " ✗"
	}
	fmt.Printf("%s  (%s)\n", r.URL, status)

	for _, ph := range singlePhaseList(r) {
		if ph.p.Present {
			fmt.Printf("  %s :  %8.2f ms\n", ph.label, ms(ph.p.Duration))
		}
	}
	fmt.Println("  " + strings.Repeat("─", 29))
	if r.Total.Present {
		fmt.Printf("  %s :  %8.2f ms\n", "Total         ", ms(r.Total.Duration))
	}
	if r.Err != nil {
		fmt.Printf("  ✗ %s: %s\n", r.FailPhase, unwrapMsg(r.Err))
	}
}

func printAggregate(a probe.Aggregate, failures []failItem, fail bool) {
	status := fmt.Sprintf("%d", a.StatusCode)
	if fail && a.StatusCode >= 400 {
		status += " ✗"
	}

	header := fmt.Sprintf("%s  (%s, %d samples", a.URL, status, a.Count)
	if len(failures) > 0 {
		header += fmt.Sprintf(", %d failed", len(failures))
	}
	header += ")"
	fmt.Println(header)

	if a.Total.Present {
		fmt.Printf("  %-14s    %9s  %9s  %9s\n", "", "min", "avg", "max")
		for _, ph := range aggPhaseList(a) {
			if ph.p.Present {
				fmt.Printf("  %s :  %6.2f ms  %6.2f ms  %6.2f ms\n",
					ph.label, ms(ph.p.Min), ms(ph.p.Avg), ms(ph.p.Max))
			}
		}
		fmt.Println("  " + strings.Repeat("─", 49))
		fmt.Printf("  %s :  %6.2f ms  %6.2f ms  %6.2f ms\n",
			"Total         ", ms(a.Total.Min), ms(a.Total.Avg), ms(a.Total.Max))
	}

	printFailureSummary(failures)
}

func printFailureSummary(failures []failItem) {
	if len(failures) == 0 {
		return
	}
	// group by (phase, message)
	type key struct{ phase, msg string }
	counts := map[key]int{}
	order := []key{}
	for _, f := range failures {
		k := key{f.phase, f.message}
		if counts[k] == 0 {
			order = append(order, k)
		}
		counts[k]++
	}
	for _, k := range order {
		n := counts[k]
		if n == 1 {
			fmt.Printf("  ✗ %s: %s\n", k.phase, k.msg)
		} else {
			fmt.Printf("  ✗ %d × %s: %s\n", n, k.phase, k.msg)
		}
	}
}

func singlePhaseList(r probe.Result) []struct {
	label string
	p     probe.Phase
} {
	return []struct {
		label string
		p     probe.Phase
	}{
		{"DNS lookup    ", r.DNS},
		{"TCP connect   ", r.Connect},
		{"TLS handshake ", r.TLS},
		{"Server (TTFB) ", r.TTFB},
		{"Transfer      ", r.Transfer},
	}
}

func aggPhaseList(a probe.Aggregate) []struct {
	label string
	p     probe.PhaseStats
} {
	return []struct {
		label string
		p     probe.PhaseStats
	}{
		{"DNS lookup    ", a.DNS},
		{"TCP connect   ", a.Connect},
		{"TLS handshake ", a.TLS},
		{"Server (TTFB) ", a.TTFB},
		{"Transfer      ", a.Transfer},
	}
}

func ms(d time.Duration) float64 {
	return float64(d.Microseconds()) / 1000
}

// ── JSON output ───────────────────────────────────────────────────────────────

type jsonPhase struct {
	MinMS float64 `json:"min_ms"`
	AvgMS float64 `json:"avg_ms"`
	MaxMS float64 `json:"max_ms"`
}

type jsonError struct {
	Phase   string `json:"phase"`
	Count   int    `json:"count"`
	Message string `json:"message"`
}

type jsonEntry struct {
	URL       string               `json:"url"`
	Status    int                  `json:"status"`
	Succeeded int                  `json:"succeeded"`
	Failed    int                  `json:"failed"`
	Phases    map[string]jsonPhase `json:"phases,omitempty"`
	Errors    []jsonError          `json:"errors,omitempty"`
}

func buildJSONEntry(rawURL string, succeeded []probe.Result, failures []failItem) jsonEntry {
	e := jsonEntry{
		URL:       rawURL,
		Succeeded: len(succeeded),
		Failed:    len(failures),
	}

	if len(succeeded) > 0 {
		a := probe.Summarize(succeeded)
		e.Status = a.StatusCode
		e.Phases = make(map[string]jsonPhase)
		addJSONPhase := func(name string, s probe.PhaseStats) {
			if s.Present {
				e.Phases[name] = jsonPhase{MinMS: ms(s.Min), AvgMS: ms(s.Avg), MaxMS: ms(s.Max)}
			}
		}
		addJSONPhase("dns", a.DNS)
		addJSONPhase("connect", a.Connect)
		addJSONPhase("tls", a.TLS)
		addJSONPhase("ttfb", a.TTFB)
		addJSONPhase("transfer", a.Transfer)
		addJSONPhase("total", a.Total)
	}

	// group failures
	type key struct{ phase, msg string }
	counts := map[key]int{}
	order := []key{}
	for _, f := range failures {
		k := key{f.phase, f.message}
		if counts[k] == 0 {
			order = append(order, k)
		}
		counts[k]++
	}
	for _, k := range order {
		e.Errors = append(e.Errors, jsonError{Phase: k.phase, Count: counts[k], Message: k.msg})
	}

	return e
}