package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"os"
	"time"
)

var (
	url        = "http://localhost:13337"
	fromText   = "/gsam2/image/maskfromtext"
	fromBboxes = "/gsam2/image/maskfrombboxes"
	fromPoints = "/gsam2/image/maskfrompoints"

	image = "truck.jpg"
)

type (
	Point struct {
		X       int  `json:"x"`
		Y       int  `json:"y"`
		Include bool `json:"include"`
	}

	BoundingBox struct {
		UpperLeft  [2]int `json:"upper_left"`
		LowerRight [2]int `json:"lower_right"`
	}

	MaskFromTextRequest struct {
		Image string `json:"image"`
		Text  string `json:"text"`
	}

	MaskFromBBoxRequest struct {
		Image  string        `json:"image"`
		Bboxes []BoundingBox `json:"bboxes"`
	}

	MaskFromPointsRequest struct {
		Image  string  `json:"image"`
		Points []Point `json:"points"`
	}

	MaskResult struct {
		Mask  string      `json:"mask"`
		Score float64     `json:"score"`
		BBox  BoundingBox `json:"bbox"`
		// fields, only populated in responses for MaskFromTextRequests
		ClassName    string      `json:"class_name"`
		DinoBBox     BoundingBox `json:"dino_bbox"`
		CenterOfMass [2]float64  `json:"center_of_mass"`
	}

	MaskResponse struct {
		Masks []MaskResult `json:"masks"`
		Image string       `json:"image"`
	}
)

func main() {
	// ensure the out directory exists
	os.Mkdir("out", 0755)

	// load the sample image and base64 encode it
	dat, err := os.ReadFile(image)
	if err != nil {
		fmt.Println(err)
		return
	}

	// post to the different endpoints
	c := &http.Client{Timeout: time.Minute}
	encImage := base64.StdEncoding.EncodeToString(dat)

	// from text
	err = doFromText(c, "fromtext", encImage, "truck. tire.")
	if err != nil {
		fmt.Printf("error %s", err)
		return
	}

	// from bboxes
	err = doFromBboxes(c, "frombboxes", encImage, []BoundingBox{
		{
			UpperLeft:  [2]int{75, 275},
			LowerRight: [2]int{1725, 850},
		},
		{
			UpperLeft:  [2]int{425, 600},
			LowerRight: [2]int{700, 875},
		},
		{
			UpperLeft:  [2]int{1375, 550},
			LowerRight: [2]int{1650, 800},
		},
		{
			UpperLeft:  [2]int{1240, 675},
			LowerRight: [2]int{1400, 750},
		},
	})
	if err != nil {
		fmt.Printf("error %s", err)
		return
	}

	// from points
	err = doFromPoints(c, "frompoints", encImage, []Point{
		{X: 500, Y: 375, Include: true},
		{X: 1125, Y: 625, Include: true},
		{X: 575, Y: 750, Include: false},
	})
	if err != nil {
		fmt.Printf("error %s", err)
		return
	}
}

func do(c *http.Client, req *http.Request, outname string) error {

	dump, err := httputil.DumpRequest(req, false)
	if err != nil {
		return err
	}
	fmt.Println("request: ", string(dump))

	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		dump, err := httputil.DumpResponse(resp, true)
		if err != nil {
			return err
		}
		fmt.Println("response: ", string(dump))
		defer resp.Body.Close()
	} else {
		dump, err := httputil.DumpResponse(resp, false)
		if err != nil {
			return err
		}
		fmt.Println("response: ", string(dump))
	}

	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	maskResp := MaskResponse{}
	err = json.Unmarshal(bodyBytes, &maskResp)
	if err != nil {
		return err
	}

	// write the masks to a file
	for _, mask := range maskResp.Masks {
		dec, err := base64.StdEncoding.DecodeString(mask.Mask)
		if err != nil {
			return err
		}
		class := ""
		if mask.ClassName != "" {
			class = "-" + mask.ClassName
		}
		os.WriteFile(fmt.Sprintf("out/%s%s-%.4f.jpg", outname, class, mask.Score), dec, 0644)
	}

	dec, err := base64.StdEncoding.DecodeString(maskResp.Image)
	if err != nil {
		return err
	}
	os.WriteFile(fmt.Sprintf("out/%s.jpg", outname), dec, 0644)

	return nil
}

func doFromText(c *http.Client, outname string, encImage string, text string) error {
	req, err := http.NewRequest("POST", url+fromText, nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept", `application/json`)

	body := MaskFromTextRequest{
		Image: encImage,
		Text:  text,
	}
	jsonBody, err := json.Marshal(body)
	if err != nil {
		return err
	}
	req.Body = io.NopCloser(bytes.NewBuffer(jsonBody))

	return do(c, req, outname)
}

func doFromBboxes(c *http.Client, outname string, encImage string, bboxes []BoundingBox) error {
	req, err := http.NewRequest("POST", url+fromBboxes, nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept", `application/json`)

	body := MaskFromBBoxRequest{
		Image:  encImage,
		Bboxes: bboxes,
	}
	jsonBody, err := json.Marshal(body)
	if err != nil {
		return err
	}
	req.Body = io.NopCloser(bytes.NewBuffer(jsonBody))

	return do(c, req, outname)
}

func doFromPoints(c *http.Client, outname string, encImage string, points []Point) error {
	req, err := http.NewRequest("POST", url+fromPoints, nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept", `application/json`)

	body := MaskFromPointsRequest{
		Image:  encImage,
		Points: points,
	}
	jsonBody, err := json.Marshal(body)
	if err != nil {
		return err
	}
	req.Body = io.NopCloser(bytes.NewBuffer(jsonBody))

	return do(c, req, outname)
}