package main

import (
    "crypto/tls"
    "flag"
    "fmt"
    "io/ioutil"
    "net/http"
    "os"
    "strings"
    "time"
)

// MeasureTTFB measures the Time to First Byte for the given URL.
func MeasureTTFB(url string, cookie string, verbose bool) (int64, error) {
    // Create a custom HTTP transport to allow measuring the TTFB.
    transport := &http.Transport{
        TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
    }

    client := &http.Client{
        Transport: transport,
    }

    // Create a new request.
    req, err := http.NewRequest("GET", url, nil)
    if err != nil {
        return 0, err
    }

    // If a cookie is provided, set it in the request header.
    if cookie != "" {
        req.Header.Set("Cookie", cookie)
        if verbose {
            fmt.Println("Using cookie:", cookie)
        }
    }

    if verbose {
        fmt.Println("Sending request to:", url)
    }

    // Record the start time.
    start := time.Now()

    // Perform the request.
    resp, err := client.Do(req)
    if err != nil {
        return 0, err
    }
    defer resp.Body.Close()

    // Measure the time taken to receive the first byte.
    ttfb := time.Since(start).Milliseconds()

    if verbose {
        fmt.Println("Received response status:", resp.Status)
    }

    return ttfb, nil
}

func main() {
    // Define and parse command-line flags.
    verbose := flag.Bool("v", false, "Enable verbose output")
    cookieFile := flag.String("c", "", "Path to file containing the authentication cookie")
    flag.Parse()

    // Check if the URL is provided as an argument.
    if flag.NArg() != 1 {
        fmt.Println("Usage: go run main.go [-v] [-c cookieFile] <URL>")
        os.Exit(1)
    }

    url := flag.Arg(0)

    // Read the cookie from the specified file if provided.
    var cookie string
    if *cookieFile != "" {
        data, err := ioutil.ReadFile(*cookieFile)
        if err != nil {
            fmt.Printf("Error reading cookie file: %v\n", err)
            os.Exit(1)
        }
        cookie = strings.TrimSpace(string(data))
    }

    // Measure the TTFB.
    ttfb, err := MeasureTTFB(url, cookie, *verbose)
    if err != nil {
        fmt.Printf("Error measuring TTFB: %v\n", err)
        os.Exit(1)
    }

    // Only print the TTFB in milliseconds by default.
    if *verbose {
        fmt.Printf("Time to First Byte for %s: %d ms\n", url, ttfb)
    } else {
        fmt.Printf("%d\n", ttfb)
    }
}