package main

import (
	"context"
	"log/slog"
	"net/http"
	"os"

	"github.com/mathiasbq/supervisor/internal/auth"
	"github.com/mathiasbq/supervisor/internal/config"
	iexec "github.com/mathiasbq/supervisor/internal/exec"
	"github.com/mathiasbq/supervisor/internal/mcp"
	"github.com/mathiasbq/supervisor/internal/registry"
	"github.com/mathiasbq/supervisor/internal/skills/brain"
	"github.com/mathiasbq/supervisor/internal/skills/org"
	"github.com/mathiasbq/supervisor/internal/skills/retrospective"
	skilldebug "github.com/mathiasbq/supervisor/internal/skills/debug"
	"github.com/mathiasbq/supervisor/internal/skills/review"
	"github.com/mathiasbq/supervisor/internal/skills/spec"
	"github.com/mathiasbq/supervisor/internal/skills/trainer"
	"github.com/mathiasbq/supervisor/internal/skills/sessionlog"
	"github.com/mathiasbq/supervisor/internal/skills/tdd"
	"github.com/mathiasbq/supervisor/internal/tier"
)

func main() {
	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))

	cfg, err := config.Load()
	if err != nil {
		logger.Error("load config", "err", err)
		os.Exit(1)
	}

	models, err := config.LoadModels(cfg.ModelsFile)
	if err != nil {
		logger.Error("load models", "err", err)
		os.Exit(1)
	}

	protocolsPrompt, err := os.ReadFile(cfg.ConfigDir + "/protocols.md")
	if err != nil {
		logger.Error("read protocols.md", "path", cfg.ConfigDir+"/protocols.md", "err", err)
		os.Exit(1)
	}

	// prependProtocols prepends the shared protocols to a skill discipline file.
	prependProtocols := func(skillPrompt []byte) string {
		return string(protocolsPrompt) + "\n---\n\n" + string(skillPrompt)
	}

	tddPrompt, err := os.ReadFile(cfg.ConfigDir + "/tdd.md")
	if err != nil {
		logger.Error("read tdd.md", "path", cfg.ConfigDir+"/tdd.md", "err", err)
		os.Exit(1)
	}

	retroPrompt, err := os.ReadFile(cfg.ConfigDir + "/retrospective.md")
	if err != nil {
		logger.Error("read retrospective.md", "path", cfg.ConfigDir+"/retrospective.md", "err", err)
		os.Exit(1)
	}

	reviewPrompt, err := os.ReadFile(cfg.ConfigDir + "/review.md")
	if err != nil {
		logger.Error("read review.md", "path", cfg.ConfigDir+"/review.md", "err", err)
		os.Exit(1)
	}

	debugPrompt, err := os.ReadFile(cfg.ConfigDir + "/debug.md")
	if err != nil {
		logger.Error("read debug.md", "path", cfg.ConfigDir+"/debug.md", "err", err)
		os.Exit(1)
	}

	specPrompt, err := os.ReadFile(cfg.ConfigDir + "/spec.md")
	if err != nil {
		logger.Error("read spec.md", "path", cfg.ConfigDir+"/spec.md", "err", err)
		os.Exit(1)
	}

	trainerReaderPrompt, err := os.ReadFile(cfg.ConfigDir + "/trainer-reader.md")
	if err != nil {
		logger.Error("read trainer-reader.md", "path", cfg.ConfigDir+"/trainer-reader.md", "err", err)
		os.Exit(1)
	}
	trainerWriterPrompt, err := os.ReadFile(cfg.ConfigDir + "/trainer-writer.md")
	if err != nil {
		logger.Error("read trainer-writer.md", "path", cfg.ConfigDir+"/trainer-writer.md", "err", err)
		os.Exit(1)
	}

	litellm := iexec.NewLiteLLM(cfg.LiteLLMBaseURL, cfg.LiteLLMAPIKey, 0)

	tierFn := func(ctx context.Context) tier.Info {
		return tier.Detect(ctx, "https://api.anthropic.com", cfg.LiteLLMBaseURL)
	}

	reg := registry.New()
	reg.Register(tdd.New(tdd.Config{
		SkillPrompt:   prependProtocols(tddPrompt),
		DefaultModel:  models.ModelFor("tdd", ""),
		CompleteFunc:  litellm.Complete,
		SessionsDir:   cfg.SessionsDir,
		IngestBaseURL: cfg.IngestBaseURL,
	}))
	reg.Register(brain.New(brain.Config{
		IngestBaseURL:  cfg.IngestBaseURL,
		IngestSvcURL:   cfg.IngestSvcURL,
		KBRetrievalURL: cfg.KBRetrievalURL,
	}))
	reg.Register(org.New(org.Config{
		TierFn: tierFn,
	}))
	reg.Register(sessionlog.New(sessionlog.Config{
		SessionsDir: cfg.SessionsDir,
	}))
	reg.Register(retrospective.New(retrospective.Config{
		SkillPrompt:  prependProtocols(retroPrompt),
		DefaultModel: models.ModelFor("retrospective", ""),
		SessionsDir:  cfg.SessionsDir,
		CompleteFunc: litellm.Complete,
	}))
	reg.Register(review.New(review.Config{
		SkillPrompt:   prependProtocols(reviewPrompt),
		DefaultModel:  models.ModelFor("review", ""),
		CompleteFunc:  litellm.Complete,
		SessionsDir:   cfg.SessionsDir,
		IngestBaseURL: cfg.IngestBaseURL,
	}))
	reg.Register(skilldebug.New(skilldebug.Config{
		SkillPrompt:   prependProtocols(debugPrompt),
		DefaultModel:  models.ModelFor("debug", ""),
		CompleteFunc:  litellm.Complete,
		SessionsDir:   cfg.SessionsDir,
		IngestBaseURL: cfg.IngestBaseURL,
	}))
	reg.Register(spec.New(spec.Config{
		SkillPrompt:   prependProtocols(specPrompt),
		DefaultModel:  models.ModelFor("spec", ""),
		CompleteFunc:  litellm.Complete,
		SessionsDir:   cfg.SessionsDir,
		IngestBaseURL: cfg.IngestBaseURL,
	}))
	reg.Register(trainer.New(trainer.Config{
		ReaderPrompt: prependProtocols(trainerReaderPrompt),
		WriterPrompt: prependProtocols(trainerWriterPrompt),
		DefaultModel: models.ModelFor("trainer", ""),
		CompleteFunc: litellm.Complete,
		SessionsDir:  cfg.SessionsDir,
		BrainDir:     cfg.BrainDir,
	}))

	var validator *auth.Validator
	if dexURL := os.Getenv("DEX_ISSUER_URL"); dexURL != "" {
		audience := os.Getenv("MCP_AUDIENCE")
		v, err := auth.NewValidator(dexURL, audience)
		if err != nil {
			logger.Error("build jwt validator", "err", err)
			os.Exit(1)
		}
		validator = v
		logger.Info("jwt auth enabled", "issuer", dexURL)
	}

	srv := mcp.NewServer(reg, cfg.MCPAuthToken, validator)
	mux := http.NewServeMux()
	mux.Handle("/mcp", srv)

	if dexURL := os.Getenv("DEX_ISSUER_URL"); dexURL != "" {
		resourceURL := os.Getenv("MCP_RESOURCE_URL")
		mux.HandleFunc("GET /.well-known/oauth-protected-resource",
			auth.ProtectedResourceHandler(resourceURL, dexURL))
	}

	addr := ":" + cfg.Port
	logger.Info("supervisor starting", "addr", addr, "version", "v0.5.0")
	if err := http.ListenAndServe(addr, mux); err != nil {
		logger.Error("server stopped", "err", err)
		os.Exit(1)
	}
}