package main

import (
	"context"
	"encoding/json"
	"math"
	"net/http"
	"strconv"
	"strings"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// ── Chat message document ────────────────────────────────────────────

type chatMessageDoc struct {
	ID        primitive.ObjectID `bson:"_id,omitempty" json:"id"`
	ProblemID string             `bson:"problemId"     json:"problemId"`
	UserID    string             `bson:"userId"        json:"userId"`
	Username  string             `bson:"username"      json:"username"`
	Message   string             `bson:"message"       json:"message"`
	CreatedAt time.Time          `bson:"createdAt"     json:"createdAt"`
}

// ── Create indexes ───────────────────────────────────────────────────

func ensureChatIndexes(coll *mongo.Collection) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Fast lookup by problemId, sorted by newest first
	_, _ = coll.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{{Key: "problemId", Value: 1}, {Key: "createdAt", Value: -1}},
	})
}

// ── GET /api/chat?problemId=...&page=1&limit=30 ─────────────────────

func handleGetChatMessages(coll *mongo.Collection) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			jsonError(w, http.StatusMethodNotAllowed, "Method not allowed")
			return
		}

		problemID := r.URL.Query().Get("problemId")
		if problemID == "" {
			jsonError(w, http.StatusBadRequest, "problemId is required")
			return
		}

		page, _ := strconv.Atoi(r.URL.Query().Get("page"))
		if page < 1 {
			page = 1
		}
		limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
		if limit < 1 || limit > 100 {
			limit = 30
		}
		skip := int64((page - 1) * limit)

		ctx := r.Context()
		filter := bson.M{"problemId": problemID}

		total, err := coll.CountDocuments(ctx, filter)
		if err != nil {
			jsonError(w, http.StatusInternalServerError, "Failed to count messages")
			return
		}
		totalPages := int(math.Ceil(float64(total) / float64(limit)))

		findOpts := options.Find().
			SetSort(bson.D{{Key: "createdAt", Value: -1}}).
			SetSkip(skip).
			SetLimit(int64(limit))

		cursor, err := coll.Find(ctx, filter, findOpts)
		if err != nil {
			jsonError(w, http.StatusInternalServerError, "Failed to fetch messages")
			return
		}
		defer cursor.Close(ctx)

		var messages []chatMessageDoc
		if err := cursor.All(ctx, &messages); err != nil {
			jsonError(w, http.StatusInternalServerError, "Failed to parse messages")
			return
		}
		if messages == nil {
			messages = []chatMessageDoc{}
		}

		jsonWrite(w, http.StatusOK, map[string]any{
			"messages":    messages,
			"totalPages":  totalPages,
			"currentPage": page,
			"total":       total,
		})
	}
}

// ── POST /api/chat ───────────────────────────────────────────────────

func handlePostChatMessage(coll *mongo.Collection, usersColl *mongo.Collection) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			jsonError(w, http.StatusMethodNotAllowed, "Method not allowed")
			return
		}

		var body struct {
			ProblemID string `json:"problemId"`
			UserID    string `json:"userId"`
			Username  string `json:"username"`
			Message   string `json:"message"`
		}
		if err := json.NewDecoder(maxBody(w, r.Body)).Decode(&body); err != nil {
			jsonError(w, http.StatusBadRequest, "Invalid JSON body")
			return
		}

		// Try JWT first, fall back to body userId for local-demo users
		userID, jwtErr := extractUserIDFromJWT(r)
		if jwtErr != nil {
			userID = strings.TrimSpace(body.UserID)
		}
		if userID == "" {
			jsonError(w, http.StatusUnauthorized, "Not authenticated")
			return
		}

		if strings.TrimSpace(body.ProblemID) == "" || strings.TrimSpace(body.Message) == "" {
			jsonError(w, http.StatusBadRequest, "problemId and message are required")
			return
		}

		// Determine username
		username := strings.TrimSpace(body.Username)
		if jwtErr == nil {
			oid, err := primitive.ObjectIDFromHex(userID)
			if err == nil {
				var uDoc userDoc
				if e := usersColl.FindOne(r.Context(), bson.M{"_id": oid}).Decode(&uDoc); e == nil {
					username = uDoc.DisplayName
				}
			}
		}
		if username == "" {
			username = "User"
		}

		// Limit message length to 500 chars
		msg := strings.TrimSpace(body.Message)
		if len(msg) > 500 {
			msg = msg[:500]
		}

		doc := chatMessageDoc{
			ProblemID: body.ProblemID,
			UserID:    userID,
			Username:  username,
			Message:   msg,
			CreatedAt: time.Now().UTC(),
		}

		result, err := coll.InsertOne(r.Context(), doc)
		if err != nil {
			jsonError(w, http.StatusInternalServerError, "Failed to save message")
			return
		}

		doc.ID = result.InsertedID.(primitive.ObjectID)

		jsonWrite(w, http.StatusCreated, doc)
	}
}

// ── DELETE /api/chat?id=...&userId=... ──────────────────────────────

func handleDeleteChatMessage(coll *mongo.Collection) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodDelete {
			jsonError(w, http.StatusMethodNotAllowed, "Method not allowed")
			return
		}

		userID, jwtErr := extractUserIDFromJWT(r)
		if jwtErr != nil {
			userID = strings.TrimSpace(r.URL.Query().Get("userId"))
		}
		if userID == "" {
			jsonError(w, http.StatusUnauthorized, "Not authenticated")
			return
		}

		msgIDHex := r.URL.Query().Get("id")
		msgOID, err := primitive.ObjectIDFromHex(msgIDHex)
		if err != nil {
			jsonError(w, http.StatusBadRequest, "Invalid message ID")
			return
		}

		ctx := r.Context()

		// Ownership check
		var doc chatMessageDoc
		if err := coll.FindOne(ctx, bson.M{"_id": msgOID}).Decode(&doc); err != nil {
			jsonError(w, http.StatusNotFound, "Message not found")
			return
		}

		if doc.UserID != userID {
			jsonError(w, http.StatusForbidden, "You can only delete your own messages")
			return
		}

		_, err = coll.DeleteOne(ctx, bson.M{"_id": msgOID})
		if err != nil {
			jsonError(w, http.StatusInternalServerError, "Failed to delete message")
			return
		}

		jsonWrite(w, http.StatusOK, map[string]any{"success": true})
	}
}