package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/joho/godotenv"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"github.com/qdrant/go-client/qdrant"
	"go.uber.org/zap"
)

var version = "0.1.0"

func main() {
	godotenv.Load()

	logger, err := zap.NewProduction()
	if err != nil {
		panic(fmt.Sprintf("failed to init logger: %v", err))
	}
	defer logger.Sync()

	qdrantURL := os.Getenv("QDRANT_URL")
	openAIKey := requireEnv(logger, "OPENAI_API_KEY")
	token := requireEnv(logger, "MCP_TOKEN")
	embeddingModel := os.Getenv("EMBEDDING_MODEL")
	embeddingBaseURL := os.Getenv("EMBEDDING_BASE_URL")
	port := envOrDefault("PORT", "8080")

	qdrantPort := 6334
	if p, err := strconv.Atoi(envOrDefault("QDRANT_PORT", "6334")); err == nil {
		qdrantPort = p
	}
	qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
	qdrantTLS := os.Getenv("QDRANT_TLS") == "true"

	var qdrantStr *qdrantStore
	if qdrantURL != "" {
		qdrantClient, err := qdrant.NewClient(&qdrant.Config{
			Host:   qdrantURL,
			Port:   qdrantPort,
			APIKey: qdrantAPIKey,
			UseTLS: qdrantTLS,
		})
		if err != nil {
			logger.Fatal("failed to create qdrant client", zap.Error(err))
		}
		defer qdrantClient.Close()

		emb := newEmbedder(openAIKey, embeddingModel, embeddingBaseURL)
		qdrantStr = newQdrantStore(qdrantClient, emb)

		if err := qdrantStr.ensureCollection(context.Background()); err != nil {
			logger.Fatal("failed to ensure collection", zap.Error(err))
		}
	} else {
		logger.Warn("QDRANT_URL not set, using SQLite FTS5 for search")
	}

	sqlitePath := envOrDefault("SQLITE_PATH", "flashcard.db")
	sqliteStore, err := newSqliteStore(sqlitePath)
	if err != nil {
		logger.Fatal("failed to init sqlite", zap.Error(err))
	}
	defer sqliteStore.close()

	mcpServer := server.NewMCPServer(
		"flashcard-mcp",
		version,
		server.WithToolCapabilities(true),
		server.WithLogging(),
	)

	registerStoreFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
	registerSearchFlashcard(mcpServer, sqliteStore, qdrantStr)
	registerUpdateFlashcard(mcpServer, sqliteStore, qdrantStr, logger)
	registerDeleteFlashcard(mcpServer, sqliteStore, qdrantStr, logger)

	streamableServer := server.NewStreamableHTTPServer(mcpServer)

	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{
			"build_date": os.Getenv("BUILD_DATE"),
			"git_commit": os.Getenv("GIT_COMMIT"),
		})
	})
	mux.Handle("/mcp", authMiddleware(token, streamableServer))

	httpServer := &http.Server{
		Addr:    ":" + port,
		Handler: corsMiddleware(mux),
	}

	logger.Info("starting", zap.String("port", port))

	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
		sig := <-sigChan
		logger.Info("shutting down", zap.String("signal", sig.String()))
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if err := streamableServer.Shutdown(shutdownCtx); err != nil {
			logger.Error("shutdown error", zap.Error(err))
		}
		sqliteStore.close()
		os.Exit(0)
	}()

	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logger.Fatal("server stopped", zap.Error(err))
	}
}

func registerStoreFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("store_flashcard",
			mcp.WithDescription("Store a flashcard for later retrieval. Use this to persist useful information discovered during a session."),
			mcp.WithString("note",
				mcp.Required(),
				mcp.Description("The markdown body — keep it concise, one paragraph usually"),
			),
			mcp.WithNumber("public_id",
				mcp.Description("Optional public identifier. Auto-generated if omitted."),
			),
			mcp.WithArray("repos",
				mcp.Description("Optional list of repository names"),
			),
			mcp.WithString("host",
				mcp.Description("Optional hostname or null if universal"),
			),
			mcp.WithArray("tags",
				mcp.Description("Optional list of tags"),
			),
			mcp.WithString("user",
				mcp.Description("Optional user identifier"),
			),
			mcp.WithString("context",
				mcp.Description("Optional context — what triggered this"),
			),
		),
		handleStoreFlashcard(sqlite, qdrant, logger),
	)
}

func registerSearchFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore) {
	mcpServer.AddTool(
		mcp.NewTool("search_flashcard",
			mcp.WithDescription("Search previously stored flashcards using semantic search. Returns the most relevant entries for the given query."),
			mcp.WithString("query",
				mcp.Required(),
				mcp.Description("The search query"),
			),
			mcp.WithNumber("limit",
				mcp.Description("Maximum number of results to return (default: 5)"),
			),
		),
		handleSearchFlashcard(sqlite, qdrant),
	)
}

func registerUpdateFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("update_flashcard",
			mcp.WithDescription("Update an existing flashcard. Creates a new version while preserving history."),
			mcp.WithNumber("public_id",
				mcp.Required(),
				mcp.Description("Public identifier of the flashcard to update"),
			),
			mcp.WithString("note",
				mcp.Description("The markdown body"),
			),
			mcp.WithArray("repos",
				mcp.Description("Optional list of repository names"),
			),
			mcp.WithString("host",
				mcp.Description("Optional hostname or null if universal"),
			),
			mcp.WithArray("tags",
				mcp.Description("Optional list of tags"),
			),
			mcp.WithString("user",
				mcp.Description("Optional user identifier"),
			),
			mcp.WithString("context",
				mcp.Description("Optional context — what triggered this"),
			),
		),
		handleUpdateFlashcard(sqlite, qdrant, logger),
	)
}

func handleUpdateFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, ok := request.Params.Arguments.(map[string]any)
		if !ok {
			return mcp.NewToolResultError("invalid arguments"), nil
		}

		publicID := parseInt64(args["public_id"])
		if publicID == 0 {
			return mcp.NewToolResultError("public_id is required and must be a positive integer"), nil
		}

		latest, err := sqlite.getLatestByPublicID(ctx, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to query: %v", err)), nil
		}
		if latest == nil {
			return mcp.NewToolResultError(fmt.Sprintf("no flashcard found with public_id %d", publicID)), nil
		}

		note := latest.Note
		repos := latest.Repos
		host := latest.Host
		tags := latest.Tags
		user := latest.User
		context_ := latest.Context

		if v := parseOptionalString(args["note"]); v != "" {
			note = v
		}
		if v := parseStringSlice(args["repos"]); v != nil {
			repos = v
		}
		if v := parseOptionalString(args["host"]); v != "" {
			host = v
		}
		if v := parseStringSlice(args["tags"]); v != nil {
			tags = v
		}
		if v := parseOptionalString(args["user"]); v != "" {
			user = v
		}
		if v := parseOptionalString(args["context"]); v != "" {
			context_ = v
		}

		version := latest.Version + 1
		id, err := sqlite.save(ctx, note, host, user, context_, repos, tags, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to persist update: %v", err)), nil
		}

		if err := sqlite.deleteFTS(ctx, latest.ID); err != nil {
			logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err))
		}
		if err := sqlite.indexFTS(ctx, id, note); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("fts index failed: %v", err)), nil
		}
		if qdrant != nil {
			if err := qdrant.index(ctx, id, publicID, version, note); err != nil {
				logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err))
			}
		}

		entry, err := sqlite.getByID(ctx, id)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to retrieve: %v", err)), nil
		}

		output, err := json.MarshalIndent(entry, "", "  ")
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to marshal result: %v", err)), nil
		}

		return mcp.NewToolResultText(string(output)), nil
	}
}

func registerDeleteFlashcard(mcpServer *server.MCPServer, sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) {
	mcpServer.AddTool(
		mcp.NewTool("delete_flashcard",
			mcp.WithDescription("Delete a flashcard. Creates a tombstone version and removes it from the vector index."),
			mcp.WithNumber("public_id",
				mcp.Required(),
				mcp.Description("Public identifier of the flashcard to delete"),
			),
		),
		handleDeleteFlashcard(sqlite, qdrant, logger),
	)
}

func handleDeleteFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, ok := request.Params.Arguments.(map[string]any)
		if !ok {
			return mcp.NewToolResultError("invalid arguments"), nil
		}

		publicID := parseInt64(args["public_id"])
		if publicID == 0 {
			return mcp.NewToolResultError("public_id is required and must be a positive integer"), nil
		}

		latest, err := sqlite.getLatestByPublicID(ctx, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to query: %v", err)), nil
		}
		if latest == nil {
			return mcp.NewToolResultError(fmt.Sprintf("no flashcard found with public_id %d", publicID)), nil
		}
		if latest.DeleteTimestamp != "" {
			return mcp.NewToolResultError(fmt.Sprintf("flashcard %d is already deleted", publicID)), nil
		}

		if _, err := sqlite.tombstone(ctx, publicID, latest.Version+1); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to tombstone: %v", err)), nil
		}

		if err := sqlite.deleteFTS(ctx, latest.ID); err != nil {
			logger.Warn("fts delete failed", zap.Int64("id", latest.ID), zap.Error(err))
		}
		if qdrant != nil {
			if err := qdrant.delete(ctx, publicID); err != nil {
				logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err))
			}
		}

		return mcp.NewToolResultText(fmt.Sprintf("flashcard %d deleted", publicID)), nil
	}
}

func handleStoreFlashcard(sqlite *sqliteStore, qdrant *qdrantStore, logger *zap.Logger) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, ok := request.Params.Arguments.(map[string]any)
		if !ok {
			return mcp.NewToolResultError("invalid arguments"), nil
		}

		note, ok := args["note"].(string)
		if !ok || note == "" {
			return mcp.NewToolResultError("note is required and must be a non-empty string"), nil
		}

		repos := parseStringSlice(args["repos"])
		host := parseOptionalString(args["host"])
		tags := parseStringSlice(args["tags"])
		user := parseOptionalString(args["user"])
		context_ := parseOptionalString(args["context"])
		publicID := parseInt64(args["public_id"])

		id, err := sqlite.save(ctx, note, host, user, context_, repos, tags, publicID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to persist: %v", err)), nil
		}

		if publicID == 0 {
			publicID = id
		}
		if err := sqlite.indexFTS(ctx, id, note); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("fts index failed: %v", err)), nil
		}
		if qdrant != nil {
			if err := qdrant.index(ctx, id, publicID, 1, note); err != nil {
				logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err))
			}
		}

		entry, err := sqlite.getByID(ctx, id)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to retrieve: %v", err)), nil
		}

		output, err := json.MarshalIndent(entry, "", "  ")
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to marshal result: %v", err)), nil
		}

		return mcp.NewToolResultText(string(output)), nil
	}
}

func handleSearchFlashcard(sqlite *sqliteStore, qdrant *qdrantStore) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, ok := request.Params.Arguments.(map[string]any)
		if !ok {
			return mcp.NewToolResultError("invalid arguments"), nil
		}

		query, ok := args["query"].(string)
		if !ok || query == "" {
			return mcp.NewToolResultError("query is required and must be a non-empty string"), nil
		}

		limit := uint64(5)
		if rawLimit := args["limit"]; rawLimit != nil {
			switch v := rawLimit.(type) {
			case float64:
				limit = uint64(v)
			case int:
				limit = uint64(v)
			case int64:
				limit = uint64(v)
			}
		}

		var results []searchResult
		var err error
		if qdrant != nil {
			results, err = qdrant.search(ctx, query, limit)
			if err != nil {
				return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
			}
			for i, r := range results {
				if cached, err := sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil {
					results[i].Entry = *cached
				}
			}
		} else {
			results, err = sqlite.searchFTS(ctx, query, limit)
			if err != nil {
				return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err)), nil
			}
		}

		output, err := json.MarshalIndent(results, "", "  ")
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to marshal results: %v", err)), nil
		}

		return mcp.NewToolResultText(string(output)), nil
	}
}

func authMiddleware(token string, next http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		authHeader := strings.ToLower(r.Header.Get("Authorization"))
		if !strings.HasPrefix(authHeader, "bearer ") || strings.TrimPrefix(authHeader, "bearer ") != token {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte(`{"error":"unauthorized"}`))
			return
		}

		next.ServeHTTP(w, r)
	}
}

func corsMiddleware(next http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusOK)
			return
		}

		next.ServeHTTP(w, r)
	}
}

func requireEnv(logger *zap.Logger, key string) string {
	val := os.Getenv(key)
	if val == "" {
		logger.Fatal("missing required env var", zap.String("key", key))
	}
	return val
}

func envOrDefault(key, defaultVal string) string {
	val := os.Getenv(key)
	if val == "" {
		return defaultVal
	}
	return val
}