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 != "" {
        req.Header.Set("Cookie", cookie)
    }

    // Create a channel to capture the time to first byte
    ttfbChan := make(chan int64)
    errorChan := make(chan error)
    statusCodeChan := make(chan int)

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

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

        // Measure the time taken to receive the first byte
        ttfb := time.Since(start).Milliseconds()
        ttfbChan <- ttfb
        statusCodeChan <- resp.StatusCode

        // Read the body to ensure we handle the response properly
        _, _ = ioutil.ReadAll(resp.Body)
    }()

    select {
    case ttfb := <-ttfbChan:
        statusCode := <-statusCodeChan
        return ttfb, statusCode, nil
    case err := <-errorChan:
        return 0, 0, err
    }
}

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