package playwrightcigo

import (
        "context"
        "fmt"
        "sync"

        "github.com/playwright-community/playwright-go"
)

var mutexBrowser sync.Mutex

type browser struct {
        instanceOf   string
        instancePort int
        uri          string
        cancel       context.CancelFunc
        count        int
}

func (b *browser) connect() (playwright.Browser, error) {
        mutexBrowser.Lock()
        defer mutexBrowser.Unlock()

        if b.count > 0 {
                b.count++
                return connect(b.instanceOf, b.uri)
        }

        if browsers == nil {
                return nil, fmt.Errorf("container is not running")
        }

        uri, cancel, err := browsers.Exec(b.instanceOf, b.instancePort)
        if err != nil {
                return nil, fmt.Errorf("could not exec chromium: %w", err)
        }
        b.uri = uri
        b.cancel = func() {
                mutexBrowser.Lock()
                defer mutexBrowser.Unlock()

                b.count--
                if b.count == 0 {
                        cancel()
                }
        }
        b.count++

        return connect(b.instanceOf, b.uri)
}

func connect(instanceOf, uri string) (playwright.Browser, error) {
        switch instanceOf {
        case "chromium":
                return pw.Chromium.Connect(uri)
        case "firefox":
                return pw.Firefox.Connect(uri)
        case "webkit":
                return pw.WebKit.Connect(uri)
        default:
                return nil, fmt.Errorf("unknown browser instance: %s", instanceOf)
        }
}

var chromium = browser{
        instanceOf:   "chromium",
        instancePort: 1024 + 3,
}

// Chromium launches a Chromium browser instance in the container
// and returns a browser object that can be used to create pages,
// navigate to websites, and perform browser automation.
//
// The connection to the browser is established via WebSockets.
// You should call browser.Close() when you're done with the browser.
// The API of the returned browser object is the Playwright API.
func Chromium() (playwright.Browser, error) {
        return chromium.connect()
}

var firefox = browser{
        instanceOf:   "firefox",
        instancePort: 1024 + 1,
}

// Firefox launches a Firefox browser instance in the container
// and returns a browser object that can be used to create pages,
// navigate to websites, and perform browser automation.
//
// The connection to the browser is established via WebSockets.
// You should call browser.Close() when you're done with the browser.
// The API of the returned browser object is the Playwright API.
func Firefox() (playwright.Browser, error) {
        return firefox.connect()
}

var webkit = browser{
        instanceOf:   "webkit",
        instancePort: 1024 + 2,
}

// Webkit launches a WebKit browser instance in the container
// and returns a browser object that can be used to create pages,
// navigate to websites, and perform browser automation.
//
// The connection to the browser is established via WebSockets.
// You should call browser.Close() when you're done with the browser.
// The API of the returned browser object is the Playwright API.
func Webkit() (playwright.Browser, error) {
        return webkit.connect()
}

package playwrightcigo

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log"
        "os/exec"
        "runtime/debug"
        "strconv"
        "strings"
        "time"

        "github.com/docker/go-connections/nat"
        "github.com/testcontainers/testcontainers-go"
        "github.com/testcontainers/testcontainers-go/wait"
)

type config struct {
        ctx        context.Context
        timeout    time.Duration
        sleeping   time.Duration
        repository string
        tag        string
        retry      int
        verbose    bool
}

type container struct {
        context    context.Context
        proxy      string
        proxyPort  int
        proxyClose func()
        browsers   testcontainers.Container
        terminate  func()
}

type module struct {
        Path    string
        Version string
        Main    bool
}

func new(version string, opts ...Option) (*container, error) {
        c := &config{
                timeout:    5 * time.Minute,
                sleeping:   200 * time.Millisecond,
                retry:      15,
                ctx:        context.Background(),
                repository: "ghcr.io/mountain-reverie/playwright-ci-go",
                tag:        "",
                verbose:    false,
        }
        for _, opt := range opts {
                opt.apply(c)
        }

        if c.tag == "" {
                tag, err := noTagVersion(version, c.verbose)
                if err != nil {
                        return nil, err
                }
                c.tag = tag
        }

        ctx, cancel := context.WithTimeout(c.ctx, c.timeout)

        timeoutSecond := int(c.timeout.Seconds())

        proxy, proxyPort, close := transparentProxy(c.retry, c.sleeping, c.verbose)

        if c.verbose {
                log.Println("Starting browser container", fmt.Sprintf("%s:%s", c.repository, c.tag))
        }
        genericContainerReq := testcontainers.GenericContainerRequest{
                ContainerRequest: testcontainers.ContainerRequest{
                        Image:           fmt.Sprintf("%s:%s", c.repository, c.tag),
                        HostAccessPorts: []int{int(proxyPort)},
                        WorkingDir:      "/src",
                        ExposedPorts:    []string{"1025/tcp", "1026/tcp", "1027/tcp"},
                        Cmd:             []string{fmt.Sprintf("sleep %v", timeoutSecond+10)},
                        WaitingFor:      wait.ForExec([]string{"echo", "ready"}),
                },
                Started: true,
        }

        browsers, err := testcontainers.GenericContainer(ctx, genericContainerReq)
        if err != nil {
                cancel()
                return nil, fmt.Errorf("could not start browser container: %w", err)
        }

        return &container{
                context:    ctx,
                proxy:      proxy,
                proxyPort:  proxyPort,
                proxyClose: close,
                browsers:   browsers,
                terminate:  cancel,
        }, nil
}

// Close terminates the container and cleans up associated resources.
func (c *container) Close() error {
        if err := c.browsers.Terminate(context.Background()); err != nil {
                return fmt.Errorf("could not terminate browser container: %w", err)
        }
        c.proxyClose()
        c.terminate()
        return nil
}

// Exec executes a browser command in the container and returns a WebSocket connection URL.
// The browser parameter should be one of: "chromium", "firefox", or "webkit".
// It also returns a cancel function to terminate the browser session.
func (c *container) Exec(browser string, containerPort int) (string, context.CancelFunc, error) {
        execCtx, execCancel := context.WithCancel(c.context)
        go func() {
                code, stuff, err := c.browsers.Exec(execCtx, []string{"node", browser + ".js", c.proxy, strconv.Itoa(c.proxyPort)})

                // Check that the context is not expired
                select {
                case <-execCtx.Done():
                        return
                default:
                }

                if err != nil {
                        log.Fatalf("Could not exec in browser container: %s", err)
                }
                if code != 0 {
                        s, err := io.ReadAll(stuff)
                        if err != nil {
                                fmt.Println("Could not read stdout/stderr from browser container:", err)
                        } else {
                                fmt.Println("Browser container output:", string(s))
                        }
                        log.Fatalf("Exec failed in browser container: %d", code)
                }
        }()

        host, err := c.browsers.Host(c.context)
        if err != nil {
                execCancel()
                return "", nil, fmt.Errorf("could not get browser host: %w", err)
        }

        p, err := port(c.context, c.browsers, host, containerPort)
        if err != nil {
                execCancel()
                return "", nil, fmt.Errorf("could not get %s port: %w", browser, err)
        }

        return fmt.Sprintf("ws://%s:%d/"+browser, host, p), execCancel, nil
}

func port(ctx context.Context, container testcontainers.Container, host string, port int) (int, error) {
        p, err := container.MappedPort(ctx, nat.Port(fmt.Sprintf("%d/tcp", port)))
        if err != nil {
                return 0, fmt.Errorf("could not get browser port: %w", err)
        }
        if err := Wait4Port(fmt.Sprintf("http://%s:%d", host, p.Int())); err != nil {
                return 0, fmt.Errorf("timeout, could not connect to browser container: %w", err)
        }
        return p.Int(), nil
}

func noTagVersion(version string, verbose bool) (string, error) {
        splitted := strings.Split(version, ".")
        if len(splitted) != 3 {
                return "", fmt.Errorf("invalid version format: %s", version)
        }

        imageVersion := fmt.Sprintf("v0.%s%02s.0", splitted[1], splitted[2])
        if imageVersion == "v0.5101.0" {
                // Workaround our CI having failed publishing this version
                imageVersion = "v0.5101.1"
        }

        found := false

        found, imageVersion = getPlaywrightCIGoFromBuildInfo(imageVersion, verbose)
        if !found {
                _, imageVersion = getPlaywrightCIGoFromGoList(imageVersion, verbose)
        }
        return imageVersion, nil
}

func filterVersion(version string) string {
        parts := strings.Split(version, "-")
        if len(parts) > 0 {
                return parts[0]
        }

        return version
}

func getPlaywrightCIGoGitVersion(imageVersion string, verbose bool) (bool, string) {
        cmd := exec.Command("git", "describe", "--tags")
        output, err := cmd.Output()
        if err != nil {
                if verbose {
                        log.Printf("could not get git version: %v\n", err)
                }
                return false, imageVersion
        }
        imageVersion = filterVersion(strings.TrimSpace(string(output)))
        if verbose {
                log.Println("Using version from git:", imageVersion)
        }
        return true, imageVersion
}

func getPlaywrightCIGoFromBuildInfo(imageVersion string, verbose bool) (bool, string) {
        if info, ok := debug.ReadBuildInfo(); !ok {
                for _, deps := range info.Deps {
                        if strings.Contains(deps.Path, "github.com/mountain-reverie/playwright-ci-go") {
                                if len(deps.Version) > 0 && deps.Version[0] == 'v' {
                                        if verbose {
                                                log.Println("Using version from build info:", deps.Version)
                                        }
                                        return true, deps.Version
                                }
                        }
                }
        }

        return false, imageVersion
}

func getPlaywrightCIGoFromGoList(imageVersion string, verbose bool) (bool, string) {
        cmd := exec.Command("go", "list", "-json", "-m", "all")
        output, err := cmd.StdoutPipe()
        if err != nil {
                if verbose {
                        log.Printf("could not get stdout pipe: %v\n", err)
                }
                return false, imageVersion
        }

        if err := cmd.Start(); err != nil {
                if verbose {
                        log.Printf("could not start command: %v\n", err)
                }
                return false, imageVersion
        }
        defer func() {
                _ = cmd.Wait()
        }()

        return parseGoListJSONStream(output, imageVersion, verbose)
}

func parseGoListJSONStream(output io.Reader, imageVersion string, verbose bool) (bool, string) {
        decoder := json.NewDecoder(output)

        defer func() {
                // Consume the rest of the stream
                if _, err := io.Copy(io.Discard, output); err != nil && verbose {
                        log.Printf("could not discard remaining output: %v\n", err)
                }
        }()

        modules := map[string]module{}

        for {
                var mod module
                if err := decoder.Decode(&mod); err != nil {
                        if err == io.EOF {
                                break
                        }
                        if verbose {
                                log.Printf("could not decode module: %v\n", err)
                        }
                        return false, imageVersion
                }

                if !strings.Contains(mod.Path, "playwright") {
                        continue
                }

                if verbose {
                        if mod.Main {
                                log.Printf("Found main module: %s\n", mod.Path)
                        } else {
                                log.Printf("Found module: %s Version: %s\n", mod.Path, mod.Version)
                        }
                }

                modules[mod.Path] = mod
        }

        version := ""
        if mod, exists := modules["github.com/playwright-community/playwright-go"]; exists {
                if verbose {
                        log.Println("Found playwright-go module:", mod.Path, "Version:", mod.Version)
                }
                if len(mod.Version) > 0 && mod.Version[0] == 'v' {
                        if verbose {
                                log.Println("Found playwright-go version:", mod.Version)
                        }
                        version = mod.Version
                }
        }

        if mod, exists := modules["github.com/mountain-reverie/playwright-ci-go"]; exists {
                if verbose {
                        log.Println("Found playwright-ci-go module:", mod.Path, "Version:", mod.Version)
                }

                if mod.Main {
                        if version != "" {
                                if verbose {
                                        log.Println("Using version from playwright-go module:", version)
                                }
                                return true, version
                        }
                        return getPlaywrightCIGoGitVersion(imageVersion, verbose)
                } else if len(mod.Version) > 0 && mod.Version[0] == 'v' {
                        if verbose {
                                log.Println("Using version from go list:", mod.Version)
                        }
                        return true, mod.Version
                } else {
                        if verbose {
                                log.Println("No version found in go list for playwright-ci-go module")
                        }
                        if version != "" {
                                if verbose {
                                        log.Println("Using version from playwright-go module:", version)
                                }
                                return false, version
                        } else {
                                return false, imageVersion
                        }
                }
        }

        if version != "" {
                if verbose {
                        log.Println("Using version from playwright-go module:", version)
                }
                return false, version
        }

        if verbose {
                log.Println("No build, module or git info found. Keeping version as:", imageVersion)
        }
        return false, imageVersion
}

// Package playwrightcigo provides a containerized solution for running
// Playwright browser tests in Go with consistent behavior across environments.
package playwrightcigo

import (
        "fmt"
        "log"
        "sync"

        "github.com/containerd/errdefs"
        "github.com/playwright-community/playwright-go"
)

var pw *playwright.Playwright
var browsers *container
var count = 0

var mutex sync.Mutex

// Install sets up the containerized Playwright environment.
// It installs the Playwright driver and creates a container to run the browsers.
// This function should be called once before using any browser.
// Options can be provided to customize the installation behavior.
// Multiple calls to Install are supported, but only the first one will
// perform the actual installation.
// Generally you want to call this in your TestMain function.
func Install(opts ...Option) error {
        mutex.Lock()
        defer mutex.Unlock()

        if count > 0 {
                count++
                return nil
        }

        c := config{}
        for _, opt := range opts {
                opt.apply(&c)
        }

        driver, err := playwright.NewDriver(&playwright.RunOptions{SkipInstallBrowsers: true, Verbose: c.verbose})
        if err != nil {
                return fmt.Errorf("error while setting up driver: %w", err)
        }

        if err := driver.Install(); err != nil {
                return fmt.Errorf("error while installing driver: %w", err)
        }

        pw, err = playwright.Run()
        if err != nil {
                return fmt.Errorf("error while starting to run playwright: %w", err)
        }

        browsers, err = new(driver.Version, opts...)
        if err != nil {
                return err
        }

        count++
        return nil
}

// Uninstall cleans up resources created by Install.
// This function should be called when you're finished with all browser testing.
// There should be as many calls to Uninstall as there were calls to Install.
// Generally you want to call this in your TestMain function.
func Uninstall() error {
        mutex.Lock()
        defer mutex.Unlock()
        count--
        if count > 0 {
                return nil
        }

        for chromium.count > 0 {
                chromium.cancel()
        }
        for firefox.count > 0 {
                firefox.cancel()
        }
        for webkit.count > 0 {
                webkit.cancel()
        }

        if err := browsers.Close(); err != nil {
                // Ignore "not found" errors since the container may have already been terminated due to timeout or other cleanup.
                if !errdefs.IsNotFound(err) {
                        return fmt.Errorf("could not close container: %w", err)
                }
                log.Println("container already closed or not found, ignoring error:", err)
        }

        if err := pw.Stop(); err != nil {
                return fmt.Errorf("could not stop playwright: %w", err)
        }
        return nil
}

package playwrightcigo

import (
        "context"
        "time"
)

// Option is an interface for configuring the behavior of playwright-ci-go functions.
type Option interface {
        apply(*config)
}

var _ Option = (*optionFunc)(nil)

type optionFunc func(*config)

func (f optionFunc) apply(c *config) {
        f(c)
}

// WithContext provides a context for cancellation and timeout control.
// This option allows you to control when operations should be canceled.
func WithContext(ctx context.Context) Option {
        return optionFunc(func(c *config) {
                c.ctx = ctx
        })
}

// WithTimeout sets the timeout of the container.
// The default timeout is 5 minutes.
func WithTimeout(timeout time.Duration) Option {
        return optionFunc(func(c *config) {
                c.timeout = timeout
        })
}

// WithRetry configures the number of retry attempts when setting up the proxy.
// The default is 15 retries.
func WithRetry(count int) Option {
        return optionFunc(func(c *config) {
                if count > 0 {
                        c.retry = count
                }
        })
}

// WithSleeping sets the sleep duration between retry attempts when setting up the proxy.
// The default is 200 milliseconds.
func WithSleeping(sleeping time.Duration) Option {
        return optionFunc(func(c *config) {
                if sleeping > 0 {
                        c.sleeping = sleeping
                }
        })
}

// WithRepository sets a custom container repository and tag.
// This is useful for using custom container images or specific versions.
// The default repository is "ghcr.io/mountain-reverie/playwright-ci-go".
func WithRepository(repository, tag string) Option {
        return optionFunc(func(c *config) {
                if repository != "" {
                        c.repository = repository
                }
                if tag != "" {
                        c.tag = tag
                }
        })
}

func WithVerbose() Option {
        return optionFunc(func(c *config) {
                c.verbose = true
        })
}

package playwrightcigo

import (
        "context"
        "errors"
        "fmt"
        "log"
        "net"
        "net/http"
        "strconv"
        "time"

        "github.com/elazarl/goproxy"
        "github.com/testcontainers/testcontainers-go"
)

func transparentProxy(retry int, sleeping time.Duration, verbose bool) (string, int, func()) {
        // Listen for incoming connections
        l, err := net.Listen("tcp", "127.0.0.1:0")
        if err != nil {
                log.Fatal("Error listening:", err)
        }

        proxy := goproxy.NewProxyHttpServer()
        proxy.Verbose = verbose

        srv := &http.Server{
                Handler:           proxy,
                ReadHeaderTimeout: time.Second * 5, // Set a reasonable ReadHeaderTimeout value
        }

        go func() {
                err := srv.Serve(l)
                if err != nil && !errors.Is(err, http.ErrServerClosed) {
                        log.Fatalf("Error serving proxy: %v", err)
                }
        }()

        _, portStr, err := net.SplitHostPort(l.Addr().String())
        if err != nil {
                log.Fatalf("Failed to parse address %s: %v", l.Addr().String(), err)
        }
        port, err := strconv.ParseInt(portStr, 10, 64)
        if err != nil {
                log.Fatalf("Failed to parse port number from address %s: %v", l.Addr().String(), err)
        }
        // Ensure the port number is within the valid range for a 16-bit unsigned integer
        if port < 0 || port > 65535 {
                log.Fatalf("Parsed port number %d is out of valid range (0-65535)", port)
        }
        if err := Wait4Port("http://"+l.Addr().String(), WithRetry(retry), WithSleeping(sleeping)); err != nil {
                log.Fatalf("Could not connect to proxy: %s", err)
        }

        return "http://" + testcontainers.HostInternal + ":" + portStr, int(port), func() {
                _ = srv.Shutdown(context.Background())
                _ = l.Close()
        }
}

// Wait4Port checks if a network service is available at the given address.
// It retries according to the provided options.
// This is useful for ensuring that servers are ready before connecting to them.
//
// Parameters:
//   - addr: The URL to check (e.g. "http://localhost:8080")
//   - opts: Configuration options for retries and timeout
func Wait4Port(addr string, opts ...Option) error {
        c := &config{
                sleeping: 200 * time.Millisecond,
                retry:    15,
                verbose:  false,
                ctx:      context.Background(),
        }
        for _, opt := range opts {
                opt.apply(c)
        }

        if err := SleepWithContext(c.ctx, c.sleeping); err != nil {
                return err
        }
        for i := 0; i < c.retry; i++ {
                req, err := http.NewRequestWithContext(c.ctx, "GET", addr, nil)
                if err != nil {
                        return fmt.Errorf("could not create a request to %s: %w", addr, err)
                }

                resp, err := http.DefaultClient.Do(req)
                if err != nil {
                        if c.verbose {
                                log.Println("could not connect to", addr, "yet, error", err, "retrying in", c.sleeping)
                        }
                        if err := SleepWithContext(c.ctx, c.sleeping); err != nil {
                                return err
                        }
                        continue
                }
                if err := resp.Body.Close(); err != nil {
                        return fmt.Errorf("could not close response body: %w", err)
                }
                return nil
        }
        return fmt.Errorf("could not connect to %s after retry and timeout", addr)
}

// SleepWithContext sleeps for the specified duration or until the context is canceled.
// It returns nil if the sleep completes or the context's error if canceled early.
func SleepWithContext(ctx context.Context, d time.Duration) error {
        select {
        case <-time.After(d):
                return nil
        case <-ctx.Done():
                return ctx.Err()
        }
}