main.go

package main

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

	"clarabase/config"
	"clarabase/middleware"
	"clarabase/routers"
	auth_service "clarabase/services/auth_service"
	"clarabase/services/cache"
	"clarabase/services/database_manager"
	feedback_service "clarabase/services/feedback"
	"clarabase/services/genai"

	resource_instances "clarabase/services/hosted_resource_instances"
	invitations_service "clarabase/services/invitations"
	organizations_service "clarabase/services/organizations"
	resource_definitions "clarabase/services/resource_definitions"
	"clarabase/services/specgen"
	supertokens_init "clarabase/services/supertokens"
	users_service "clarabase/services/users"

	"github.com/golang-migrate/migrate/v4"
	_ "github.com/golang-migrate/migrate/v4/database/postgres"
	_ "github.com/golang-migrate/migrate/v4/source/file"

	chiMiddleware "github.com/go-chi/chi/middleware"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/cors"
	"github.com/go-chi/httprate"
	httprateredis "github.com/go-chi/httprate-redis"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/joho/godotenv"
	"github.com/supertokens/supertokens-golang/recipe/session"
	"github.com/supertokens/supertokens-golang/supertokens"
)

// TODO:
// - user deletion flow: remove from all tenants, organization; if last user, delete organization and all data

func init() {
	godotenv.Load()
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile) // Add file:line to logs
	config := config.Config()

	m, err := migrate.New("file://db/migrations", config.DBConnectionURI)
	if err != nil {
		log.Fatalf("Could not create migration: %v", err)
	}
	if err := m.Up(); err != nil && err != migrate.ErrNoChange {
		log.Fatalf("Could not run migrations: %v", err)
	}

	db, err := pgxpool.New(context.Background(), config.DBConnectionURI)
	if err != nil {
		log.Fatalf("Could not connect to database: %s --- %s", config.DBConnectionURI, os.Getenv("POSTGRESQL_CONNECTION_URI"))
	}
	defer db.Close()

	err = supertokens_init.InitializeSuperTokens(db, config)
	if err != nil {
		panic(err.Error())
	}

	r := chi.NewRouter()
	r.Use(cors.Handler(cors.Options{
		AllowedOrigins:   []string{config.AppURL},
		AllowedMethods:   []string{"GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"},
		AllowedHeaders:   append([]string{"Accept", "Authorization", "Content-Type", "X-CSRF-Token", "Clarabase-UI-Origin", "clarabase-ui-origin"}, supertokens.GetAllCORSHeaders()...),
		AllowCredentials: true,
	}))
	r.Use(chiMiddleware.RequestID)
	r.Use(chiMiddleware.Logger)
	r.Use(chiMiddleware.Recoverer)
	r.Use(chiMiddleware.Heartbeat("/healthz"))
	r.Use(supertokens.Middleware)
	r.Use(middleware.WithSecurityHeaders)
	r.Use(httprate.Limit(
		100,
		time.Minute,
		httprate.WithKeyByIP(),
		httprate.WithResponseHeaders(httprate.ResponseHeaders{
			Limit:      "X-RateLimit-Limit",
			Remaining:  "X-RateLimit-Remaining",
			Reset:      "X-RateLimit-Reset",
			RetryAfter: "Retry-After",
		}),
		httprateredis.WithRedisLimitCounter(&httprateredis.Config{
			Client: cache.RedisClient(),
		}),
	))

	// Tenant-facing routes
	// /auth
	// /api
	r.Use(routers.InitializeTenantHostRouter(db))

	r.Get("/sessioninfo", session.VerifySession(nil, sessioninfo))

	r.Route("/public", func(r chi.Router) {
		r.Use(chiMiddleware.Timeout(30 * time.Second))
		r.Use(middleware.AggressiveRateLimit)
		r.Use(middleware.PublicEndpointSecurity)

		r.Route("/feedback", func(r chi.Router) {
			r.Use(middleware.PublicFeedbackRateLimit)
			feedback_service.RegisterRoutes(r)
		})
	})

	// Routes just for the frontend app
	r.Route("/app", func(r chi.Router) {
		r.Use(middleware.WithTransaction(db))
		r.Use(middleware.AuthenticateViaAppSession)

		r.Group(func(r chi.Router) {
			r.Use(chiMiddleware.Timeout(30 * time.Second))
			r.Route("/users", users_service.RegisterRoutes)
			r.Route("/organizations", organizations_service.RegisterRoutes)
			r.Route("/invitations", invitations_service.RegisterRoutes)
			r.Route("/specs", specgen.RegisterRoutes)
			r.Route("/account", auth_service.RegisterRoutes)
			r.Route("/resource-definitions", resource_definitions.RegisterRoutes)
			r.Route("/database-configs", database_manager.RegisterRoutes)
			r.Route("/feedback", feedback_service.RegisterRoutes)
			r.Post("/api-token", auth_service.HandleIssueTenantTokenFromAppSession)
			r.Get("/data-viewer/{resource}", resource_instances.ListResources)
		})

		r.Group(func(r chi.Router) {
			r.Use(chiMiddleware.Timeout(5 * time.Minute))
			r.Use(middleware.PublicGenAIRateLimit)
			r.Route("/chat", genai.RegisterRoutes)
		})
	})

	log.Println(`
=======================
|     Estoy listo!    |
=======================
	`)
	http.ListenAndServe(fmt.Sprintf(":%d", 8787), r)
}

func sessioninfo(w http.ResponseWriter, r *http.Request) {
	sessionContainer := session.GetSessionFromRequestContext(r.Context())

	if sessionContainer == nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("no session found"))
		return
	}
	sessionData, err := sessionContainer.GetSessionDataInDatabase()
	if err != nil {
		err = supertokens.ErrorHandler(err, r, w)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
		}
		return
	}
	db, err := pgxpool.New(context.Background(), config.Config().DBConnectionURI)
	if err != nil {
		log.Printf("Could not connect to database: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
	defer db.Close()
	trx, err := db.Begin(context.Background())
	if err != nil {
		log.Printf("Could not begin transaction: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
	defer trx.Rollback(context.Background())
	orgID, err := organizations_service.OrganizationMembershipByUserId(&trx, sessionContainer.GetUserID(), "owner")
	if err != nil {
		log.Printf("Could not get organization by user ID: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
	sessionContainer.MergeIntoAccessTokenPayload(map[string]interface{}{
		string(config.OrgIDKey): orgID,
	})
	w.WriteHeader(http.StatusOK)
	w.Header().Add("content-type", "application/json")
	bytes, err := json.Marshal(map[string]interface{}{
		"sessionHandle":      sessionContainer.GetHandle(),
		"userId":             sessionContainer.GetUserID(),
		"accessTokenPayload": sessionContainer.GetAccessTokenPayload(),
		"sessionData":        sessionData,
	})
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("error in converting to json"))
	} else {
		w.Write(bytes)
	}
}