gooseek/backend/cmd/labs-svc/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/cors"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gooseek/backend/internal/labs"
	"github.com/gooseek/backend/internal/llm"
	"github.com/gooseek/backend/pkg/config"
	"github.com/google/uuid"
)

type ReportStore struct {
	reports map[string]*labs.Report
}

func NewReportStore() *ReportStore {
	return &ReportStore{
		reports: make(map[string]*labs.Report),
	}
}

func (s *ReportStore) Save(report *labs.Report) {
	s.reports[report.ID] = report
}

func (s *ReportStore) Get(id string) *labs.Report {
	return s.reports[id]
}

func (s *ReportStore) List(limit, offset int) []*labs.Report {
	result := make([]*labs.Report, 0)
	i := 0
	for _, r := range s.reports {
		if i >= offset && len(result) < limit {
			result = append(result, r)
		}
		i++
	}
	return result
}

func (s *ReportStore) Delete(id string) bool {
	if _, ok := s.reports[id]; ok {
		delete(s.reports, id)
		return true
	}
	return false
}

func main() {
	cfg, err := config.Load()
	if err != nil {
		log.Fatal("Failed to load config:", err)
	}

	var llmClient llm.Client
	if cfg.OpenAIAPIKey != "" {
		client, err := llm.NewOpenAIClient(llm.ProviderConfig{
			ProviderID: "openai",
			APIKey:     cfg.OpenAIAPIKey,
			ModelKey:   "gpt-4o-mini",
		})
		if err != nil {
			log.Fatal("Failed to create OpenAI client:", err)
		}
		llmClient = client
	} else if cfg.AnthropicAPIKey != "" {
		client, err := llm.NewAnthropicClient(llm.ProviderConfig{
			ProviderID: "anthropic",
			APIKey:     cfg.AnthropicAPIKey,
			ModelKey:   "claude-3-5-sonnet-20241022",
		})
		if err != nil {
			log.Fatal("Failed to create Anthropic client:", err)
		}
		llmClient = client
	}

	generator := labs.NewGenerator(llmClient)
	store := NewReportStore()

	app := fiber.New(fiber.Config{
		BodyLimit:    100 * 1024 * 1024,
		ReadTimeout:  60 * time.Second,
		WriteTimeout: 60 * time.Second,
	})

	app.Use(logger.New())
	app.Use(cors.New())

	app.Get("/health", func(c *fiber.Ctx) error {
		return c.JSON(fiber.Map{"status": "ok"})
	})

	app.Post("/api/v1/labs/generate", func(c *fiber.Ctx) error {
		var req struct {
			Query   string      `json:"query"`
			Data    interface{} `json:"data"`
			Theme   string      `json:"theme,omitempty"`
			Locale  string      `json:"locale,omitempty"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()

		report, err := generator.GenerateReport(ctx, labs.GenerateOptions{
			Query:  req.Query,
			Data:   req.Data,
			Theme:  req.Theme,
			Locale: req.Locale,
		})
		if err != nil {
			return c.Status(500).JSON(fiber.Map{"error": err.Error()})
		}

		store.Save(report)

		return c.JSON(report)
	})

	app.Post("/api/v1/labs/visualize", func(c *fiber.Ctx) error {
		var req struct {
			Type   string      `json:"type"`
			Title  string      `json:"title"`
			Data   interface{} `json:"data"`
			Config interface{} `json:"config,omitempty"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := createVisualizationFromRequest(generator, req.Type, req.Title, req.Data, req.Config)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/chart", func(c *fiber.Ctx) error {
		var req struct {
			Type     string          `json:"type"`
			Title    string          `json:"title"`
			Labels   []string        `json:"labels"`
			Datasets []labs.ChartDataset `json:"datasets"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		var viz labs.Visualization
		switch req.Type {
		case "bar", "bar_chart":
			if len(req.Datasets) > 0 {
				viz = generator.CreateBarChart(req.Title, req.Labels, req.Datasets[0].Data)
			}
		case "line", "line_chart":
			viz = generator.CreateLineChart(req.Title, req.Labels, req.Datasets)
		case "pie", "pie_chart":
			if len(req.Datasets) > 0 {
				viz = generator.CreatePieChart(req.Title, req.Labels, req.Datasets[0].Data)
			}
		default:
			viz = generator.CreateBarChart(req.Title, req.Labels, req.Datasets[0].Data)
		}

		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/table", func(c *fiber.Ctx) error {
		var req struct {
			Title   string            `json:"title"`
			Columns []labs.TableColumn `json:"columns"`
			Rows    []labs.TableRow    `json:"rows"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateTable(req.Title, req.Columns, req.Rows)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/stat-cards", func(c *fiber.Ctx) error {
		var req struct {
			Title string           `json:"title"`
			Cards []labs.StatCard  `json:"cards"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateStatCards(req.Title, req.Cards)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/kpi", func(c *fiber.Ctx) error {
		var req struct {
			Title  string      `json:"title"`
			Value  interface{} `json:"value"`
			Change float64     `json:"change"`
			Unit   string      `json:"unit"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateKPI(req.Title, req.Value, req.Change, req.Unit)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/heatmap", func(c *fiber.Ctx) error {
		var req struct {
			Title   string      `json:"title"`
			XLabels []string    `json:"xLabels"`
			YLabels []string    `json:"yLabels"`
			Values  [][]float64 `json:"values"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateHeatmap(req.Title, req.XLabels, req.YLabels, req.Values)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/code", func(c *fiber.Ctx) error {
		var req struct {
			Title    string `json:"title"`
			Code     string `json:"code"`
			Language string `json:"language"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateCodeBlock(req.Title, req.Code, req.Language)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/markdown", func(c *fiber.Ctx) error {
		var req struct {
			Title   string `json:"title"`
			Content string `json:"content"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateMarkdown(req.Title, req.Content)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/tabs", func(c *fiber.Ctx) error {
		var req struct {
			Title string          `json:"title"`
			Tabs  []labs.TabItem  `json:"tabs"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateTabs(req.Title, req.Tabs)
		return c.JSON(viz)
	})

	app.Post("/api/v1/labs/accordion", func(c *fiber.Ctx) error {
		var req struct {
			Title string               `json:"title"`
			Items []labs.AccordionItem `json:"items"`
		}

		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request"})
		}

		viz := generator.CreateAccordion(req.Title, req.Items)
		return c.JSON(viz)
	})

	app.Get("/api/v1/labs/reports", func(c *fiber.Ctx) error {
		limit := c.QueryInt("limit", 20)
		offset := c.QueryInt("offset", 0)

		reports := store.List(limit, offset)
		return c.JSON(fiber.Map{"reports": reports, "count": len(reports)})
	})

	app.Get("/api/v1/labs/reports/:id", func(c *fiber.Ctx) error {
		id := c.Params("id")
		report := store.Get(id)
		if report == nil {
			return c.Status(404).JSON(fiber.Map{"error": "Report not found"})
		}
		return c.JSON(report)
	})

	app.Delete("/api/v1/labs/reports/:id", func(c *fiber.Ctx) error {
		id := c.Params("id")
		if store.Delete(id) {
			return c.JSON(fiber.Map{"success": true})
		}
		return c.Status(404).JSON(fiber.Map{"error": "Report not found"})
	})

	app.Post("/api/v1/labs/reports/:id/export", func(c *fiber.Ctx) error {
		id := c.Params("id")
		format := c.Query("format", "html")

		report := store.Get(id)
		if report == nil {
			return c.Status(404).JSON(fiber.Map{"error": "Report not found"})
		}

		switch format {
		case "html":
			html := exportToHTML(report)
			c.Set("Content-Type", "text/html")
			c.Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s.html\"", report.ID))
			return c.SendString(html)
		case "json":
			c.Set("Content-Type", "application/json")
			c.Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s.json\"", report.ID))
			return c.JSON(report)
		default:
			return c.Status(400).JSON(fiber.Map{"error": "Unsupported format"})
		}
	})

	port := getEnvInt("LABS_SVC_PORT", 3031)
	log.Printf("labs-svc listening on :%d", port)
	log.Fatal(app.Listen(fmt.Sprintf(":%d", port)))
}

func createVisualizationFromRequest(g *labs.Generator, vizType, title string, data, config interface{}) labs.Visualization {
	switch vizType {
	case "bar_chart":
		return parseChartRequest(g, labs.VizBarChart, title, data)
	case "line_chart":
		return parseChartRequest(g, labs.VizLineChart, title, data)
	case "pie_chart":
		return parseChartRequest(g, labs.VizPieChart, title, data)
	case "table":
		return parseTableRequest(g, title, data)
	case "stat_cards":
		return parseStatCardsRequest(g, title, data)
	case "kpi":
		return parseKPIRequest(g, title, data)
	case "markdown":
		content := ""
		if dataMap, ok := data.(map[string]interface{}); ok {
			content, _ = dataMap["content"].(string)
		}
		return g.CreateMarkdown(title, content)
	case "code_block":
		code, lang := "", ""
		if dataMap, ok := data.(map[string]interface{}); ok {
			code, _ = dataMap["code"].(string)
			lang, _ = dataMap["language"].(string)
		}
		return g.CreateCodeBlock(title, code, lang)
	default:
		return g.CreateMarkdown(title, fmt.Sprintf("%v", data))
	}
}

func parseChartRequest(g *labs.Generator, vizType labs.VisualizationType, title string, data interface{}) labs.Visualization {
	dataMap, ok := data.(map[string]interface{})
	if !ok {
		return g.CreateMarkdown(title, "Invalid chart data")
	}

	labels := make([]string, 0)
	if labelsRaw, ok := dataMap["labels"].([]interface{}); ok {
		for _, l := range labelsRaw {
			labels = append(labels, fmt.Sprintf("%v", l))
		}
	}

	values := make([]float64, 0)
	if valuesRaw, ok := dataMap["values"].([]interface{}); ok {
		for _, v := range valuesRaw {
			switch val := v.(type) {
			case float64:
				values = append(values, val)
			case int:
				values = append(values, float64(val))
			}
		}
	}

	switch vizType {
	case labs.VizBarChart:
		return g.CreateBarChart(title, labels, values)
	case labs.VizPieChart:
		return g.CreatePieChart(title, labels, values)
	default:
		return g.CreateLineChart(title, labels, []labs.ChartDataset{{Label: title, Data: values}})
	}
}

func parseTableRequest(g *labs.Generator, title string, data interface{}) labs.Visualization {
	dataMap, ok := data.(map[string]interface{})
	if !ok {
		return g.CreateMarkdown(title, "Invalid table data")
	}

	columns := make([]labs.TableColumn, 0)
	if colsRaw, ok := dataMap["columns"].([]interface{}); ok {
		for _, c := range colsRaw {
			if colMap, ok := c.(map[string]interface{}); ok {
				col := labs.TableColumn{}
				if v, ok := colMap["key"].(string); ok {
					col.Key = v
				}
				if v, ok := colMap["label"].(string); ok {
					col.Label = v
				}
				columns = append(columns, col)
			}
		}
	}

	rows := make([]labs.TableRow, 0)
	if rowsRaw, ok := dataMap["rows"].([]interface{}); ok {
		for _, r := range rowsRaw {
			if rowMap, ok := r.(map[string]interface{}); ok {
				rows = append(rows, labs.TableRow(rowMap))
			}
		}
	}

	return g.CreateTable(title, columns, rows)
}

func parseStatCardsRequest(g *labs.Generator, title string, data interface{}) labs.Visualization {
	dataMap, ok := data.(map[string]interface{})
	if !ok {
		return g.CreateMarkdown(title, "Invalid stat cards data")
	}

	cards := make([]labs.StatCard, 0)
	if cardsRaw, ok := dataMap["cards"].([]interface{}); ok {
		for _, c := range cardsRaw {
			if cardMap, ok := c.(map[string]interface{}); ok {
				card := labs.StatCard{ID: uuid.New().String()}
				if v, ok := cardMap["title"].(string); ok {
					card.Title = v
				}
				if v, ok := cardMap["value"]; ok {
					card.Value = v
				}
				if v, ok := cardMap["change"].(float64); ok {
					card.Change = v
				}
				if v, ok := cardMap["color"].(string); ok {
					card.Color = v
				}
				cards = append(cards, card)
			}
		}
	}

	return g.CreateStatCards(title, cards)
}

func parseKPIRequest(g *labs.Generator, title string, data interface{}) labs.Visualization {
	dataMap, ok := data.(map[string]interface{})
	if !ok {
		return g.CreateKPI(title, data, 0, "")
	}

	value := dataMap["value"]
	change := 0.0
	if v, ok := dataMap["change"].(float64); ok {
		change = v
	}
	unit := ""
	if v, ok := dataMap["unit"].(string); ok {
		unit = v
	}

	return g.CreateKPI(title, value, change, unit)
}

func exportToHTML(report *labs.Report) string {
	sectionsHTML := ""
	for _, section := range report.Sections {
		vizHTML := ""
		for _, viz := range section.Visualizations {
			dataJSON, _ := json.Marshal(viz.Data)
			vizHTML += fmt.Sprintf(`
				<div class="visualization" data-type="%s">
					<h3>%s</h3>
					<div class="viz-container" data-config='%s'></div>
				</div>
			`, viz.Type, viz.Title, string(dataJSON))
		}
		sectionsHTML += fmt.Sprintf(`
			<section>
				<h2>%s</h2>
				%s
			</section>
		`, section.Title, vizHTML)
	}

	return fmt.Sprintf(`<!DOCTYPE html>
<html>
<head>
	<meta charset="UTF-8">
	<title>%s</title>
	<style>
		body { font-family: system-ui, sans-serif; max-width: 1200px; margin: 0 auto; padding: 20px; }
		h1 { color: #1f2937; }
		h2 { color: #374151; border-bottom: 2px solid #e5e7eb; padding-bottom: 10px; }
		.visualization { margin: 20px 0; padding: 20px; background: #f9fafb; border-radius: 8px; }
		.viz-container { min-height: 200px; }
	</style>
	<script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
</head>
<body>
	<h1>%s</h1>
	%s
</body>
</html>`, report.Title, report.Title, sectionsHTML)
}

func getEnvInt(key string, defaultValue int) int {
	if val := os.Getenv(key); val != "" {
		var result int
		if _, err := fmt.Sscanf(val, "%d", &result); err == nil {
			return result
		}
	}
	return defaultValue
}