package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"log"
	"math"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"google.golang.org/genai"
)

func main() {
	ctx, can := signal.NotifyContext(context.Background(), syscall.SIGINT)
	defer can()
	if err := run(ctx); err != nil && ctx.Err() == nil {
		panic(err)
	}
}

func run(ctx context.Context) error {
	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	n := fs.Int("n", 1, "loops")
	model := fs.String("m", "gemini-2.5-flash", "model to use")

	if err := fs.Parse(os.Args[1:]); err != nil {
		return err
	}
	city := strings.Title(fs.Args()[0])

	client, err := genai.NewClient(ctx, nil)
	if err != nil {
		return err
	}

	descriptors := []string{
		`Literally poison in the air that politicians are adamant against doing anything to mitigate`,
		`Drought worst-case`,
		`Wildfire risk`,
		`Hurricane`,
		`Earthquake`,
		`Tornado`,
		`Crime`,
		`Consistent Air Pollution`,
		`Transient Air Pollution (from wildfires, etc)`,
		`Colorado river distribution conflict / reduced capacity`,
		`PFAS`,
		`Homelessness`,
		`Neighbors disappointed in church attendance`,
		`Hot`,
		`Personal gripes`,
		`Alcohol universal`,
		`Few traffic destination hubs`,
		`Popularity of booze`,
		`Tourist destination (eg Branson-like)`,
		`Weed universal`,
		`Greater cost of living than Utah`,
		`Traffic >= Bend at 3k/sqmi`,
		`Lots of schools`,
		`Coffee shops`,
		`Yummy resturaunts`,
		`Board game stores`,
		`Home for < 400k`,
		`20min drive to fun`,
		`Would we get out and be active? Bike++`,
		`Close airport`,
		`DnD Players`,
		`Trees!`,
		`Population is college educated`,
		`Home for < 600k`,
		`Software developers per capita`,
		`Fiber-speed internet`,
		`Organic Groceries`,
		`Close hospital for emergency situations`,
		`Home for < 800`,
		`Hematologist reasonably close`,
		`Does it have the vibes?`,
		`Concrete Jungle`,
	}

	tmpl, err := template.New("tmpl").Parse(`
For the city {{.City}}, answer the following descriptors '1' for true or '0' for false. Output only a JSON object mapping descriptors to answers.

{{ range .Descriptors -}}
* {{ . }}
{{ end }}
   `)
	if err != nil {
		return err
	}
	w := bytes.NewBuffer(nil)
	if err := tmpl.Execute(w, map[string]any{
		"City":        city,
		"Descriptors": descriptors,
	}); err != nil {
		return err
	}
	//panic(w.String())

	results := make([][]float64, len(descriptors))
	for i := 0; i < *n; i++ {
		result, err := client.Models.GenerateContent(
			ctx,
			*model,
			genai.Text(w.String()),
			nil,
		)
		if err != nil {
			return err
		}
		text := "{" + strings.Split(strings.Split(result.Text(), "{")[1], "}")[0] + "}"

		var m map[string]any
		if err := json.Unmarshal([]byte(text), &m); err != nil {
			return err
		}
		for j, d := range descriptors {
			if v, ok := m[d]; ok {
				k := fmt.Sprint(v)
				if k == "1" {
					results[j] = append(results[j], 1.0)
				} else {
					results[j] = append(results[j], 0.0)
				}
			}
		}
	}

	log.Printf("%+v", results)
	//log.Println(text)
	for i := range descriptors {
		sum := 0.0
		for j := range results[i] {
			sum += results[i][j]
		}
		avg := sum / float64(len(results[i]))
		fmt.Println(math.Round(avg))
	}

	return ctx.Err()
}