// Package shoot photographs an application with the Chrome already on the
// machine.
//
// squint never downloads or ships a browser. An earlier tool in this family
// pinned an exact renderer in a 300MB container so that stored baseline
// images would match across machines; squint stores no baselines, so it has
// nothing to pin and borrows whatever Chrome is installed.
package shoot

import (
	"context"
	"encoding/json"
	"fmt"
	"image"
	"image/png"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"

	"git.mills.io/prologic/squint/internal/config"
	"github.com/chromedp/cdproto/network"
	cdpruntime "github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/chromedp"
)

// Browser is a running Chrome that every shot is taken through. One
// browser for a whole shoot: starting Chrome is the slowest thing squint
// does, and a shoot takes dozens of pictures.
type Browser struct {
	ctx    context.Context
	cancel []context.CancelFunc

	// sessions caches one login per user for the life of a shoot. Logging
	// in per shot is both wasteful and wrong: a matrix of thirty-odd
	// screens signed in as the same person walked straight into Parley's
	// login rate limiter and every admin screen came back a 429. An
	// application's brute-force defences are not something a screenshot
	// tool gets to opt out of.
	mu       sync.Mutex
	sessions map[string][]*http.Cookie
}

// FindChrome returns the path to an installed Chrome, or an error naming
// what was looked for. squint is useless without one and says so plainly
// rather than failing somewhere deeper.
func FindChrome() (string, error) {
	if p := os.Getenv("SQUINT_CHROME"); p != "" {
		if _, err := os.Stat(p); err == nil {
			return p, nil
		}
		return "", fmt.Errorf("SQUINT_CHROME is set to %q, which is not there", p)
	}
	var candidates []string
	switch runtime.GOOS {
	case "darwin":
		candidates = []string{
			"/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
			"/Applications/Chromium.app/Contents/MacOS/Chromium",
		}
	default:
		candidates = []string{"google-chrome", "google-chrome-stable", "chromium", "chromium-browser"}
	}
	for _, c := range candidates {
		if strings.Contains(c, "/") {
			if _, err := os.Stat(c); err == nil {
				return c, nil
			}
			continue
		}
		if p, err := exec.LookPath(c); err == nil {
			return p, nil
		}
	}
	return "", fmt.Errorf("no Chrome found (looked for %s); install one or set SQUINT_CHROME",
		strings.Join(candidates, ", "))
}

// Open starts a headless Chrome.
func Open(ctx context.Context) (*Browser, error) {
	path, err := FindChrome()
	if err != nil {
		return nil, err
	}
	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.ExecPath(path),
		chromedp.Flag("headless", "new"),
		chromedp.Flag("hide-scrollbars", true),
		chromedp.Flag("disable-gpu", true),
		// No force-device-scale-factor here: every shot sets the scale
		// explicitly through the protocol instead, so it is the config's
		// to choose and cannot be inherited from whatever display the
		// machine happens to have.
	)
	allocCtx, cancelAlloc := chromedp.NewExecAllocator(ctx, opts...)
	browserCtx, cancelBrowser := chromedp.NewContext(allocCtx)
	if err := chromedp.Run(browserCtx); err != nil {
		cancelBrowser()
		cancelAlloc()
		return nil, fmt.Errorf("starting Chrome at %s: %w", path, err)
	}
	return &Browser{
		ctx:      browserCtx,
		cancel:   []context.CancelFunc{cancelBrowser, cancelAlloc},
		sessions: map[string][]*http.Cookie{},
	}, nil
}

// Close stops Chrome.
func (b *Browser) Close() {
	for i := len(b.cancel) - 1; i >= 0; i-- {
		b.cancel[i]()
	}
}

// Shoot renders one shot and returns the image.
//
// Each shot gets its own tab, so a scenario that leaves a drawer open or a
// dialog up cannot leak into the next one.
func (b *Browser) Shoot(cfg *config.Config, base string, shot config.Shot) (image.Image, error) {
	tabCtx, cancelTab := chromedp.NewContext(b.ctx)
	defer cancelTab()
	ctx, cancel := context.WithTimeout(tabCtx, 60*time.Second)
	defer cancel()

	var cookies []*http.Cookie
	if shot.Scenario.Login != "" {
		var err error
		if cookies, err = b.session(cfg, base, shot.Scenario.Login); err != nil {
			return nil, err
		}
	}

	tasks := chromedp.Tasks{
		chromedp.EmulateViewport(int64(shot.Viewport.Width), int64(shot.Viewport.Height),
			chromedp.EmulateScale(cfg.Scale)),
		chromedp.ActionFunc(setCookies(base, cookies, shot.Theme.Cookies)),
		chromedp.Navigate(base + shot.Scenario.Path),
		chromedp.WaitReady("body", chromedp.ByQuery),
	}
	if shot.Scenario.WaitFor != "" {
		tasks = append(tasks, chromedp.WaitVisible(shot.Scenario.WaitFor, chromedp.ByQuery))
	}
	for _, st := range shot.Scenario.Steps {
		task, err := step(st)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", shot.Name, err)
		}
		tasks = append(tasks, task)
	}
	if masks := append(append([]string{}, cfg.Mask...), shot.Scenario.Mask...); len(masks) > 0 {
		tasks = append(tasks, chromedp.ActionFunc(func(ctx context.Context) error {
			js, err := maskScript(masks)
			if err != nil {
				return err
			}
			return chromedp.Evaluate(js, nil).Do(ctx)
		}))
	}
	// Stop the world before the shutter. A transition caught a few
	// milliseconds either side is pixels that differ between two runs and
	// mean nothing; killing motion removes the whole class rather than
	// tuning a tolerance against it.
	tasks = append(tasks,
		chromedp.ActionFunc(func(ctx context.Context) error {
			return chromedp.Evaluate(stopMotion, nil).Do(ctx)
		}),
		chromedp.ActionFunc(waitFonts),
	)

	var buf []byte
	tasks = append(tasks, chromedp.FullScreenshot(&buf, 100))
	if err := chromedp.Run(ctx, tasks); err != nil {
		return nil, fmt.Errorf("%s: %w", shot.Name, err)
	}
	return png.Decode(strings.NewReader(string(buf)))
}

func step(s config.Step) (chromedp.Action, error) {
	switch {
	case s.Click != "":
		return chromedp.Click(s.Click, chromedp.ByQuery), nil
	case s.Type != "":
		if s.Into == "" {
			return nil, fmt.Errorf("a type step needs an `into` selector")
		}
		return chromedp.SendKeys(s.Into, s.Type, chromedp.ByQuery), nil
	case s.WaitFor != "":
		return chromedp.WaitVisible(s.WaitFor, chromedp.ByQuery), nil
	case s.Wait > 0:
		return chromedp.Sleep(s.Wait), nil
	case s.JS != "":
		return chromedp.ActionFunc(func(ctx context.Context) error {
			return chromedp.Evaluate(s.JS, nil).Do(ctx)
		}), nil
	}
	return nil, fmt.Errorf("a step does nothing: set click, type, wait_for, wait or js")
}

// session returns the cookies for a user, logging in once per shoot and
// reusing them. Keyed by base URL as well as user, so a check -- which
// runs two application instances on two ports -- does not hand the
// baseline's session to the candidate.
func (b *Browser) session(cfg *config.Config, base, user string) ([]*http.Cookie, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	key := base + "\x00" + user
	if c, ok := b.sessions[key]; ok {
		return c, nil
	}
	c, err := login(cfg, base, user)
	if err != nil {
		return nil, err
	}
	b.sessions[key] = c
	return c, nil
}

// login drives the application's real login form with an ordinary HTTP
// client and hands the resulting cookies to the browser. Using the real
// form rather than forging a session means a broken login page is caught
// by the tool that photographs it.
func login(cfg *config.Config, base, user string) ([]*http.Cookie, error) {
	jar, _ := cookiejar.New(nil)
	client := &http.Client{
		Jar:           jar,
		Timeout:       10 * time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
	}
	form := url.Values{}
	for k, v := range cfg.Auth.Form {
		form.Set(k, config.Expand(v, map[string]string{"User": user}))
	}
	resp, err := client.PostForm(base+cfg.Auth.Path, form)
	if err != nil {
		return nil, fmt.Errorf("logging in as %s: %w", user, err)
	}
	resp.Body.Close()
	u, _ := url.Parse(base)
	cookies := jar.Cookies(u)
	if len(cookies) == 0 {
		return nil, fmt.Errorf("logging in as %s returned no cookies (status %d): check auth.form", user, resp.StatusCode)
	}
	return cookies, nil
}

// setCookies puts the session and theme cookies in place before the page
// loads, so it is drawn signed in and themed on the first and only paint.
func setCookies(base string, session []*http.Cookie, theme map[string]string) func(context.Context) error {
	u, _ := url.Parse(base)
	host := u.Hostname()
	return func(ctx context.Context) error {
		// Cookies belong to the browser, not to the tab. Without this, the
		// first shot that asked for a dark theme left its cookie behind and
		// every later shot inherited it -- the light renderings came out
		// dark, under file names that said otherwise, which is worse than
		// failing. A shot has to start from nothing.
		if err := network.ClearBrowserCookies().Do(ctx); err != nil {
			return err
		}
		set := func(name, value string) error {
			return network.SetCookie(name, value).WithDomain(host).WithPath("/").Do(ctx)
		}
		for _, c := range session {
			if err := set(c.Name, c.Value); err != nil {
				return err
			}
		}
		for k, v := range theme {
			if err := set(k, v); err != nil {
				return err
			}
		}
		return nil
	}
}

func waitFonts(ctx context.Context) error {
	var ready bool
	return chromedp.Evaluate(`document.fonts.ready.then(() => true)`, &ready,
		func(p *cdpruntime.EvaluateParams) *cdpruntime.EvaluateParams { return p.WithAwaitPromise(true) }).Do(ctx)
}

// maskScript hides the masked elements with visibility, not display, so
// they keep their space: a masked token must not shorten the row it sits
// in, or masking would itself move the page.
func maskScript(selectors []string) (string, error) {
	b, err := json.Marshal(selectors)
	if err != nil {
		return "", err
	}
	return `(() => {
  for (const sel of ` + string(b) + `) {
    for (const el of document.querySelectorAll(sel)) el.style.visibility = 'hidden';
  }
})()`, nil
}

const stopMotion = `(() => {
  const s = document.createElement('style');
  s.textContent = '*,*::before,*::after{transition:none!important;animation:none!important;caret-color:transparent!important;scroll-behavior:auto!important}';
  document.head.appendChild(s);
})()`