package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"time"

	bb "github.com/chorus-services/backbeat/internal/backbeat"
	"github.com/nats-io/nats.go"
	"gopkg.in/yaml.v3"
)

type scoreFile struct {
	Score bb.Score `yaml:"score"`
}

func main() {
	cluster := flag.String("cluster", "chorus-aus-01", "cluster id")
	agentID := flag.String("id", "bzzz-1", "agent id")
	scorePath := flag.String("score", "./configs/sample-score.yaml", "score yaml path")
	natsURL := flag.String("nats", nats.DefaultURL, "nats url")
	flag.Parse()

	buf, err := os.ReadFile(*scorePath)
	if err != nil {
		log.Fatal(err)
	}
	var s scoreFile
	if err := yaml.Unmarshal(buf, &s); err != nil {
		log.Fatal(err)
	}
	score := s.Score

	nc, err := nats.Connect(*natsURL)
	if err != nil {
		log.Fatal(err)
	}
	defer nc.Drain()

	hlc := bb.NewHLC(*agentID)
	state := "planning"
	waiting := 0
	beatsLeft := 0

	nc.Subscribe(fmt.Sprintf("backbeat.%s.beat", *cluster), func(m *nats.Msg) {
		var bf bb.BeatFrame
		if err := json.Unmarshal(m.Data, &bf); err != nil {
			return
		}
		phase, _ := bb.PhaseFor(score.Phases, int(bf.BeatIndex))
		switch phase {
		case "plan":
			state = "planning"
			beatsLeft = 0
		case "work":
			if waiting == 0 && rand.Float64() < 0.3 {
				waiting = 1
			}
			if waiting > 0 {
				state = "waiting"
				beatsLeft = score.WaitBudget.Help - waiting
				waiting++
				if waiting > score.WaitBudget.Help {
					state = "executing"
					waiting = 0
				}
			} else {
				state = "executing"
				beatsLeft = 0
			}
		case "review":
			state = "review"
			waiting = 0
			beatsLeft = 0
		}

		sc := bb.StatusClaim{
			AgentID:   *agentID,
			TaskID:    "ucxl://demo/task",
			BeatIndex: bf.BeatIndex,
			State:     state,
			WaitFor:   nil,
			BeatsLeft: beatsLeft,
			Progress:  rand.Float64(),
			Notes:     "proto",
			HLC:       hlc.Next(),
		}
		payload, _ := json.Marshal(sc)
		nc.Publish("backbeat.status."+*agentID, payload)
	})

	log.Printf("AgentSim %s started (cluster=%s)\n", *agentID, *cluster)
	for {
		time.Sleep(10 * time.Second)
	}
}