package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"net"
)

func WithSession(ctx context.Context, fn func(context.Context, *TCPTransport) error) error {
	return WithGecko(ctx, func(ctx context.Context, c *net.TCPConn) (_retErr error) {
		maybeSetErr := func(_err error) {
			if _retErr == nil && _err != nil {
				_retErr = _err
			}
		}

		hello, err := receive[MsgHello](c)
		if err != nil {
			return err
		}
		if hello.ApplicationType != "gecko" || hello.MarionetteProtocol != 3 {
			return fmt.Errorf("I only know how to speak marionette protocol v3 to gecko, but got response: %#v", hello)
		}

		t := NewTransport(c)

		resp1, err := DoCommand[any](t, "WebDriver:NewSession", map[string]bool{
			"strictFileInteractability": true,
		})
		if err != nil {
			return err
		}
		log.Printf("resp1 = %#v", resp1)

		defer func() {
			_, err := DoCommand[any](t, "Marionette:Quit", struct{}{})
			maybeSetErr(err)
			maybeSetErr(c.CloseWrite())
			_, err = io.Copy(io.Discard, c)
			maybeSetErr(err)
		}()

		return fn(ctx, t)
	})
}