package mcp

import (
	"context"
	"fmt"
	"strings"

	"github.com/mark3labs/mcp-go/mcp"
	mcpserver "github.com/mark3labs/mcp-go/server"
	"github.com/pkg/errors"

	storepb "github.com/usememos/memos/proto/gen/store"
	"github.com/usememos/memos/server/auth"
	"github.com/usememos/memos/store"
)

type attachmentJSON struct {
	Name         string `json:"name"`
	Creator      string `json:"creator"`
	CreateTime   int64  `json:"create_time"`
	Filename     string `json:"filename"`
	Type         string `json:"type"`
	Size         int64  `json:"size"`
	StorageType  string `json:"storage_type"`
	ExternalLink string `json:"external_link,omitempty"`
	Memo         string `json:"memo,omitempty"`
}

func storeAttachmentToJSON(ctx context.Context, stores *store.Store, a *store.Attachment) (attachmentJSON, error) {
	creator, err := lookupUsername(ctx, stores, a.CreatorID)
	if err != nil {
		return attachmentJSON{}, errors.Wrap(err, "lookup attachment creator username")
	}
	j := attachmentJSON{
		Name:       "attachments/" + a.UID,
		Creator:    creator,
		CreateTime: a.CreatedTs,
		Filename:   a.Filename,
		Type:       a.Type,
		Size:       a.Size,
	}
	switch a.StorageType {
	case storepb.AttachmentStorageType_LOCAL:
		j.StorageType = "LOCAL"
	case storepb.AttachmentStorageType_S3:
		j.StorageType = "S3"
		j.ExternalLink = a.Reference
	case storepb.AttachmentStorageType_EXTERNAL:
		j.StorageType = "EXTERNAL"
		j.ExternalLink = a.Reference
	default:
		j.StorageType = "DATABASE"
	}
	if a.MemoUID != nil && *a.MemoUID != "" {
		j.Memo = "memos/" + *a.MemoUID
	}
	return j, nil
}

func storeAttachmentToJSONWithUsernames(a *store.Attachment, usernamesByID map[int32]string) (attachmentJSON, error) {
	creator, err := lookupUsernameFromCache(usernamesByID, a.CreatorID)
	if err != nil {
		return attachmentJSON{}, errors.Wrap(err, "lookup attachment creator username from cache")
	}
	j := attachmentJSON{
		Name:       "attachments/" + a.UID,
		Creator:    creator,
		CreateTime: a.CreatedTs,
		Filename:   a.Filename,
		Type:       a.Type,
		Size:       a.Size,
	}
	switch a.StorageType {
	case storepb.AttachmentStorageType_LOCAL:
		j.StorageType = "LOCAL"
	case storepb.AttachmentStorageType_S3:
		j.StorageType = "S3"
		j.ExternalLink = a.Reference
	case storepb.AttachmentStorageType_EXTERNAL:
		j.StorageType = "EXTERNAL"
		j.ExternalLink = a.Reference
	default:
		j.StorageType = "DATABASE"
	}
	if a.MemoUID != nil && *a.MemoUID != "" {
		j.Memo = "memos/" + *a.MemoUID
	}
	return j, nil
}

func parseAttachmentUID(name string) (string, error) {
	uid, ok := strings.CutPrefix(name, "attachments/")
	if !ok || uid == "" {
		return "", errors.Errorf(`attachment name must be "attachments/<uid>", got %q`, name)
	}
	return uid, nil
}

func (s *MCPService) registerAttachmentTools(mcpSrv *mcpserver.MCPServer) {
	mcpSrv.AddTool(mcp.NewTool("list_attachments",
		mcp.WithDescription("List attachments owned by the authenticated user. Supports pagination and optional filtering by linked memo."),
		mcp.WithNumber("page_size", mcp.Description("Maximum attachments to return (1–100, default 20)")),
		mcp.WithNumber("page", mcp.Description("Zero-based page index (default 0)")),
		mcp.WithString("memo", mcp.Description(`Filter by linked memo resource name, e.g. "memos/abc123"`)),
	), s.handleListAttachments)

	mcpSrv.AddTool(mcp.NewTool("get_attachment",
		mcp.WithDescription("Get a single attachment's metadata by resource name. Requires authentication."),
		mcp.WithString("name", mcp.Required(), mcp.Description(`Attachment resource name, e.g. "attachments/abc123"`)),
	), s.handleGetAttachment)

	mcpSrv.AddTool(mcp.NewTool("delete_attachment",
		mcp.WithDescription("Permanently delete an attachment and its stored file. Requires authentication and ownership."),
		mcp.WithString("name", mcp.Required(), mcp.Description(`Attachment resource name, e.g. "attachments/abc123"`)),
	), s.handleDeleteAttachment)

	mcpSrv.AddTool(mcp.NewTool("link_attachment_to_memo",
		mcp.WithDescription("Link an existing attachment to a memo. Requires authentication and ownership of the attachment."),
		mcp.WithString("name", mcp.Required(), mcp.Description(`Attachment resource name, e.g. "attachments/abc123"`)),
		mcp.WithString("memo", mcp.Required(), mcp.Description(`Memo resource name, e.g. "memos/abc123"`)),
	), s.handleLinkAttachmentToMemo)
}

func (s *MCPService) handleListAttachments(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	userID, err := extractUserID(ctx)
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	pageSize := req.GetInt("page_size", 20)
	if pageSize <= 0 {
		pageSize = 20
	}
	if pageSize > 100 {
		pageSize = 100
	}
	page := req.GetInt("page", 0)
	if page < 0 {
		page = 0
	}

	limit := pageSize + 1
	offset := page * pageSize
	find := &store.FindAttachment{
		CreatorID: &userID,
		Limit:     &limit,
		Offset:    &offset,
	}

	if memoName := req.GetString("memo", ""); memoName != "" {
		memoUID, err := parseMemoUID(memoName)
		if err != nil {
			return mcp.NewToolResultError(err.Error()), nil
		}
		memo, err := s.store.GetMemo(ctx, &store.FindMemo{UID: &memoUID})
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to find memo: %v", err)), nil
		}
		if memo == nil {
			return mcp.NewToolResultError("memo not found"), nil
		}
		find.MemoID = &memo.ID
	}

	attachments, err := s.store.ListAttachments(ctx, find)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to list attachments: %v", err)), nil
	}

	hasMore := len(attachments) > pageSize
	if hasMore {
		attachments = attachments[:pageSize]
	}
	creatorIDs := make([]int32, 0, len(attachments))
	for _, attachment := range attachments {
		creatorIDs = append(creatorIDs, attachment.CreatorID)
	}
	usernamesByID, err := preloadUsernames(ctx, s.store, creatorIDs)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to preload attachment creators: %v", err)), nil
	}

	results := make([]attachmentJSON, len(attachments))
	for i, a := range attachments {
		result, err := storeAttachmentToJSONWithUsernames(a, usernamesByID)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to resolve attachment creator: %v", err)), nil
		}
		results[i] = result
	}

	type listResponse struct {
		Attachments []attachmentJSON `json:"attachments"`
		HasMore     bool             `json:"has_more"`
	}
	out, err := marshalJSON(listResponse{Attachments: results, HasMore: hasMore})
	if err != nil {
		return nil, err
	}
	return mcp.NewToolResultText(out), nil
}

func (s *MCPService) handleGetAttachment(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	userID := auth.GetUserID(ctx)

	uid, err := parseAttachmentUID(req.GetString("name", ""))
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	attachment, err := s.store.GetAttachment(ctx, &store.FindAttachment{UID: &uid})
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to get attachment: %v", err)), nil
	}
	if attachment == nil {
		return mcp.NewToolResultError("attachment not found"), nil
	}

	if err := s.checkAttachmentAccess(ctx, attachment, userID); err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	result, err := storeAttachmentToJSON(ctx, s.store, attachment)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to resolve attachment creator: %v", err)), nil
	}
	out, err := marshalJSON(result)
	if err != nil {
		return nil, err
	}
	return mcp.NewToolResultText(out), nil
}

func (s *MCPService) handleDeleteAttachment(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	userID, err := extractUserID(ctx)
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	uid, err := parseAttachmentUID(req.GetString("name", ""))
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	attachment, err := s.store.GetAttachment(ctx, &store.FindAttachment{UID: &uid, CreatorID: &userID})
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to find attachment: %v", err)), nil
	}
	if attachment == nil {
		return mcp.NewToolResultError("attachment not found"), nil
	}

	if err := s.store.DeleteAttachment(ctx, &store.DeleteAttachment{ID: attachment.ID}); err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to delete attachment: %v", err)), nil
	}
	return mcp.NewToolResultText(`{"deleted":true}`), nil
}

func (s *MCPService) handleLinkAttachmentToMemo(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	userID, err := extractUserID(ctx)
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	uid, err := parseAttachmentUID(req.GetString("name", ""))
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	attachment, err := s.store.GetAttachment(ctx, &store.FindAttachment{UID: &uid})
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to get attachment: %v", err)), nil
	}
	if attachment == nil {
		return mcp.NewToolResultError("attachment not found"), nil
	}
	if attachment.CreatorID != userID {
		return mcp.NewToolResultError("permission denied"), nil
	}

	memoUID, err := parseMemoUID(req.GetString("memo", ""))
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}
	memo, err := s.store.GetMemo(ctx, &store.FindMemo{UID: &memoUID})
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to get memo: %v", err)), nil
	}
	if memo == nil {
		return mcp.NewToolResultError("memo not found"), nil
	}
	if err := checkMemoOwnership(memo, userID); err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	if err := s.store.UpdateAttachment(ctx, &store.UpdateAttachment{
		ID:     attachment.ID,
		MemoID: &memo.ID,
	}); err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to link attachment: %v", err)), nil
	}

	// Re-fetch to get updated memo UID.
	updated, err := s.store.GetAttachment(ctx, &store.FindAttachment{ID: &attachment.ID})
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to fetch updated attachment: %v", err)), nil
	}
	result, err := storeAttachmentToJSON(ctx, s.store, updated)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to resolve attachment creator: %v", err)), nil
	}
	out, err := marshalJSON(result)
	if err != nil {
		return nil, err
	}
	return mcp.NewToolResultText(out), nil
}