package main
import (
    "os"
    "context"
    "fmt"
    "io/ioutil"
    "log"
    "time"
    "flag"
    "github.com/mafredri/cdp"
    "github.com/mafredri/cdp/devtool"
    "github.com/mafredri/cdp/protocol/page"
    "github.com/mafredri/cdp/rpcc"
)
func main() {
    folderPtr := flag.String("folder", "", "folder path for screenshots: example c:\\temp\\screens\\")
    chromePtr := flag.String("chrome", "http://localhost:9222", "chrome connection - example: http://localhost:9222")
    
    widthPtr := flag.Int("width", 1280, "screencast width")
    heightPtr := flag.Int("height", 720, "screencast height")
    qualityPtr := flag.Int("quality", 100, "screencast quality")
    
    flag.Parse()
    if err := run(*folderPtr, *chromePtr, *widthPtr, *heightPtr, *qualityPtr); err != nil {
        panic(err)
    }
}
func run(folder string, chromeConnection string, width int, height int, quality int) error {
    ctx, cancel := context.WithCancel(context.TODO())
    defer cancel()
    
    chromePath := chromeConnection
    folderPath := folder
    devt := devtool.New(chromePath)
    pageTarget, err := devt.Get(ctx, devtool.Page)
    if err != nil {
        return err
    }
    conn, err := rpcc.DialContext(ctx, pageTarget.WebSocketDebuggerURL)
    if err != nil {
        return err
    }
    defer conn.Close()
    c := cdp.NewClient(conn)
    err = c.Page.Enable(ctx)
    if err != nil {
        return err
    }
    
    screencastFrame, err := c.Page.ScreencastFrame(ctx)
    if err != nil {
        return err
    }
    go func() {
        defer screencastFrame.Close()
        for {
            ev, err := screencastFrame.Recv()
            if err != nil {
                log.Printf("Failed to receive ScreencastFrame: %v", err)
                os.Exit(0)
            }
            log.Printf("Got frame with sessionID: %d: %+v", ev.SessionID, ev.Metadata)
            err = c.Page.ScreencastFrameAck(ctx, page.NewScreencastFrameAckArgs(ev.SessionID))
            if err != nil {
                log.Printf("Failed to ack ScreencastFrame: %v", err)
                os.Exit(0)
            }
            
            name := fmt.Sprintf("screencast_frame-%d.png", ev.Metadata.Timestamp.Time().Unix())
            
            filePath := folderPath + name
            
            go func() {
                err = ioutil.WriteFile(filePath, ev.Data, 0644)
                if err != nil {
                    log.Printf("Failed to write ScreencastFrame to %q: %v", name, err)
                }
            }()
        }
    }()
    screencastArgs := page.NewStartScreencastArgs().
        SetQuality(quality).
        SetMaxWidth(width).
        SetMaxHeight(height).
        SetEveryNthFrame(1).
        SetFormat("png")
    err = c.Page.StartScreencast(ctx, screencastArgs)
    if err != nil {
        return err
    }
    
    time.Sleep(600 * time.Second)
    err = c.Page.StopScreencast(ctx)
    if err != nil {
        return err
    }
    return nil
}