package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

func value(name, fallback string) string {
	if v := os.Getenv(name); v != "" {
		return v
	}
	return fallback
}

func run() error {
	key := os.Getenv("LAB0_API_KEY")
	if key == "" {
		return fmt.Errorf("set LAB0_API_KEY to a Read & write key")
	}
	base := strings.TrimRight(value("LAB0_BASE_URL", "https://app.lab0.ai/app"), "/")
	namespace := value("LAB0_NAMESPACE", "quickstart")
	client := &http.Client{Timeout: 200 * time.Second}
	request := func(method, route, contentType string, body io.Reader) (map[string]any, error) {
		req, err := http.NewRequest(method, base+route, body)
		if err != nil {
			return nil, err
		}
		req.Header.Set("Authorization", "Bearer "+key)
		if contentType != "" {
			req.Header.Set("Content-Type", contentType)
		}
		resp, err := client.Do(req)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()
		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			text, _ := io.ReadAll(resp.Body)
			return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, text)
		}
		var result map[string]any
		err = json.NewDecoder(resp.Body).Decode(&result)
		return result, err
	}

	path := value("LAB0_FILE_PATH", "plan.md")
	document, err := os.Open(path)
	if err != nil {
		return err
	}
	defer document.Close()
	var upload bytes.Buffer
	form := multipart.NewWriter(&upload)
	part, err := form.CreateFormFile("file", filepath.Base(path))
	if err != nil {
		return err
	}
	if _, err = io.Copy(part, document); err != nil {
		return err
	}
	if err = form.WriteField("namespace", namespace); err != nil {
		return err
	}
	if err = form.Close(); err != nil {
		return err
	}
	uploaded, err := request("POST", "/api/knowledge/files", form.FormDataContentType(), &upload)
	if err != nil {
		return err
	}
	fileID, ok := uploaded["id"].(string)
	if !ok {
		return fmt.Errorf("upload response has no file id")
	}
	ready := false
	for attempt := 0; attempt < 90; attempt++ {
		file, err := request("GET", "/api/knowledge/files/"+fileID, "", nil)
		if err != nil {
			return err
		}
		if file["status"] == "ready" {
			ready = true
			break
		}
		if file["status"] == "failed" {
			return fmt.Errorf("processing failed: %v", file["error"])
		}
		time.Sleep(2 * time.Second)
	}
	if !ready {
		return fmt.Errorf("processing timed out; inspect file %s", fileID)
	}
	payload, err := json.Marshal(map[string]any{
		"prompt": "How is the rollout run?", "namespace": namespace, "file_ids": []string{fileID},
	})
	if err != nil {
		return err
	}
	answer, err := request("POST", "/api/knowledge/respond", "application/json", bytes.NewReader(payload))
	if err != nil {
		return err
	}
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	return enc.Encode(answer)
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}