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, cookies []string) (int64, int, 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, 0, err
    }

    // If cookies are provided, set them in the request header.
    if len(cookies) > 0 {
        for _, cookie := range cookies {
            if !strings.Contains(cookie, "=") {
                return 0, 0, fmt.Errorf("invalid cookie format")
            }
            req.Header.Add("Cookie", strings.TrimSpace(cookie))
        }
    }

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

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

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

    return ttfb, resp.StatusCode, nil
}

func main() {
    // Define and parse command-line flags.
    cookieFile := flag.String("c", "", "Path to file containing the authentication cookies")
    flag.Parse()

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

    url := flag.Arg(0)

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

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

    // Print the HTTP status code and TTFB.
    fmt.Printf("%d,%d\n", statusCode, ttfb)
}

