package main

import (
    "context"
    "encoding/json"
    "fmt"
    "io/ioutil"
    "log"
    "os"
    "os/exec"
    "strings"
    "time"

    "github.com/docker/docker/api/types"
    "github.com/docker/docker/api/types/events"
    "github.com/docker/docker/api/types/filters"
    "github.com/docker/docker/client"
    "github.com/robfig/cron/v3"
)

const (
    notifyScript = "/notify.sh"
)

// ContainerDiff represents the state of a container
type ContainerDiff struct {
    ID         string
    Image      string
    Labels     map[string]string
}

// getRunningContainers fetches the list of running containers with relevant labels
func getRunningContainers(cli *client.Client) ([]ContainerDiff, error) {
    containers, err := cli.ContainerList(context.Background(), types.ContainerListOptions{})
    if err != nil {
        return nil, err
    }

    var diffs []ContainerDiff
    for _, container := range containers {
        if _, ok := container.Labels["oculus.containerid"]; ok {
            diffs = append(diffs, ContainerDiff{
                ID:     container.ID,
                Image:  container.Image,
                Labels: container.Labels,
            })
        }
    }

    return diffs, nil
}

// saveDiffs writes the container diffs to a file
func saveDiffs(diffs []ContainerDiff, filePath string) error {
    data, err := json.Marshal(diffs)
    if err != nil {
        return err
    }

    return ioutil.WriteFile(filePath, data, 0644)
}

// loadDiffs reads the container diffs from a file
func loadDiffs(filePath string) ([]ContainerDiff, error) {
    data, err := ioutil.ReadFile(filePath)
    if err != nil {
        return nil, err
    }

    var diffs []ContainerDiff
    err = json.Unmarshal(data, &diffs)
    return diffs, err
}

// profileContainers logs the current state of containers to a file
func profileContainers(cli *client.Client) {
    diffs, err := getRunningContainers(cli)
    if err != nil {
        log.Fatalf("Error profiling containers: %v", err)
    }

    for _, diff := range diffs {
        if diff.Labels["oculus.mode"] == "profile" {
            filePath := fmt.Sprintf("%s.json", diff.Labels["oculus.containerid"])
            err := saveDiffs([]ContainerDiff{diff}, filePath)
            if err != nil {
                log.Fatalf("Error saving diffs: %v", err)
            }
        }
    }

    log.Println("Profiled current containers and saved to file")
}

// compareContainers compares the current state of containers with the saved state
func compareContainers(cli *client.Client) {
    currentDiffs, err := getRunningContainers(cli)
    if err != nil {
        log.Fatalf("Error getting current containers: %v", err)
    }

    for _, current := range currentDiffs {
        if current.Labels["oculus.mode"] == "monitor" {
            filePath := fmt.Sprintf("%s.json", current.Labels["oculus.containerid"])
            savedDiffs, err := loadDiffs(filePath)
            if err != nil {
                log.Printf("Error loading saved diffs for %s: %v", current.ID, err)
                continue
            }

            if len(savedDiffs) > 0 {
                saved := savedDiffs[0]
                if current.Image != saved.Image {
                    log.Printf("Container %s changed: Image %s -> %s", current.ID, saved.Image, current.Image)
                    runNotifyScript(current.ID, current.Image, saved.Image)
                }
                if current.Labels["oculus.ignorelist"] != saved.Labels["oculus.ignorelist"] {
                    log.Printf("Container %s ignore list changed: %s -> %s", current.ID, saved.Labels["oculus.ignorelist"], current.Labels["oculus.ignorelist"])
                    runNotifyScript(current.ID, current.Image, saved.Image)
                }
            } else {
                log.Printf("New container detected: ID %s, Image %s", current.ID, current.Image)
                runNotifyScript(current.ID, current.Image, "")
            }
        }
    }
}

// runNotifyScript executes the notification script with the container details
func runNotifyScript(containerID, newImage, oldImage string) {
    cmd := exec.Command(notifyScript, containerID, newImage, oldImage)
    err := cmd.Run()
    if err != nil {
        log.Printf("Error running notify script: %v", err)
    }
}

// watchDockerEvents listens for Docker events and triggers comparisons
func watchDockerEvents(cli *client.Client) {
    eventFilter := filters.NewArgs()
    eventFilter.Add("type", "container")
    eventFilter.Add("event", "start")
    eventFilter.Add("event", "stop")
    eventFilter.Add("event", "destroy")

    messages, errs := cli.Events(context.Background(), types.EventsOptions{
        Filters: eventFilter,
    })

    for {
        select {
        case event := <-messages:
            log.Printf("Docker event received: %s for container %s", event.Action, event.ID)
            compareContainers(cli)
        case err := <-errs:
            log.Printf("Error watching Docker events: %v", err)
            return
        }
    }
}

// scheduler sets up the cron jobs for profiling
func scheduler(cli *client.Client) {
    c := cron.New()
    c.AddFunc("@every 1h", func() { profileContainers(cli) })
    c.Start()
}

func main() {
    cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
    if (err != nil) {
        log.Fatalf("Error creating Docker client: %v", err)
    }

    go scheduler(cli)

    go watchDockerEvents(cli)

    // Keep the program running
    select {}
}