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) (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 a cookie is provided, set it in the request header.
    if cookie != "" {
        // Clean up and format the cookie properly
        cookie = strings.TrimSpace(cookie)
        if !strings.Contains(cookie, "=") {
            return 0, 0, fmt.Errorf("invalid cookie format")
        }
        req.Header.Set("Cookie", 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 cookie")
    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 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, statusCode, err := MeasureTTFB(url, cookie)
    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)
}