package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"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)
      --json        Output results as JSON instead of text
  -h, --help        Show this help

Examples:
  latprobe https://example.com
  latprobe -n 5 https://example.com https://www.google.com
  latprobe --json https://example.com | jq .
`

func main() {
	count := flag.Int("count", 1, "number of requests per URL")
	flag.IntVar(count, "n", 1, "number of requests per URL (shorthand)")
	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(1)
	}

	failed := false
	var jsonEntries []jsonEntry

	for i, url := range urls {
		results := collectSamples(url, *count, &failed)
		if len(results) == 0 {
			continue
		}

		if *jsonOut {
			jsonEntries = append(jsonEntries, toJSONEntry(probe.Summarize(results)))
			continue
		}

		if i > 0 {
			fmt.Println()
		}
		if *count == 1 {
			printResult(results[0])
		} else {
			printAggregate(probe.Summarize(results))
		}
	}

	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(1)
		}
	}

	if failed {
		os.Exit(1)
	}
}

func collectSamples(url string, count int, failed *bool) []probe.Result {
	results := make([]probe.Result, 0, count)
	for i := range count {
		r := probe.Measure(url)
		if r.Err != nil {
			fmt.Fprintf(os.Stderr, "error %s (sample %d/%d): %v\n", url, i+1, count, r.Err)
			*failed = true
			return nil
		}
		results = append(results, r)
	}
	return results
}

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

func printResult(r probe.Result) {
	fmt.Printf("%s  (%d)\n", r.URL, r.StatusCode)
	for _, ph := range singlePhases(r) {
		if ph.p.Present {
			fmt.Printf("  %s :  %8.2f ms\n", ph.label, ms(ph.p.Duration))
		}
	}
	fmt.Println("  " + strings.Repeat("─", 29))
	fmt.Printf("  %s :  %8.2f ms\n", "Total         ", ms(r.Total.Duration))
}

func printAggregate(a probe.Aggregate) {
	fmt.Printf("%s  (%d, %d samples)\n", a.URL, a.StatusCode, a.Count)
	fmt.Printf("  %-14s    %9s  %9s  %9s\n", "", "min", "avg", "max")
	for _, ph := range aggPhases(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))
}

func singlePhases(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 aggPhases(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 jsonEntry struct {
	URL     string               `json:"url"`
	Status  int                  `json:"status"`
	Samples int                  `json:"samples"`
	Phases  map[string]jsonPhase `json:"phases"`
}

func toJSONEntry(a probe.Aggregate) jsonEntry {
	e := jsonEntry{
		URL:     a.URL,
		Status:  a.StatusCode,
		Samples: a.Count,
		Phases:  make(map[string]jsonPhase),
	}
	add := 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),
			}
		}
	}
	add("dns", a.DNS)
	add("connect", a.Connect)
	add("tls", a.TLS)
	add("ttfb", a.TTFB)
	add("transfer", a.Transfer)
	add("total", a.Total)
	return e
}