package tools

import (
	"context"
	"fmt"

	"github.com/ClareAI/astra-insight-mcp/internal/gateway"
	"github.com/ClareAI/astra-insight-mcp/internal/mcp"
)

func registerListDocuments(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "list_knowledge_documents",
		Description: "List all documents in the agent's knowledge base. Shows document names, word counts, token counts, indexing status, and hit counts. Use this to understand what information the agent has access to and identify knowledge gaps.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"dataset_id": {Type: "string", Description: "Knowledge base (dataset) ID"},
				"page":       {Type: "integer", Description: "Page number for pagination"},
				"limit":      {Type: "integer", Description: "Max documents to return (default 100)"},
			},
			Required: []string{"dataset_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		datasetID, _ := args["dataset_id"].(string)
		if apiKey == "" || datasetID == "" {
			return nil, fmt.Errorf("api_key and dataset_id are required")
		}

		opts := gateway.DocListOpts{}
		if d, ok := args["page"].(float64); ok && d > 0 {
			opts.Page = int(d)
		}
		if d, ok := args["limit"].(float64); ok && d > 0 {
			opts.Limit = int(d)
		}

		return gw.ListDocuments(ctx, apiKey, datasetID, opts)
	})
}

func registerImportKnowledgeURL(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "import_knowledge_from_url",
		Description: "Import one or more web pages into the knowledge base. Content is extracted, chunked, and indexed. Requires a valid dataset_id UUID — call create_knowledge_base first to get one.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"dataset_id": {Type: "string", Description: "Knowledge base (dataset) UUID from create_knowledge_base"},
				"title":      {Type: "string", Description: "Title for the imported documents"},
				"urls":       {Type: "array", Description: "List of web page URLs to import (max 50)"},
			},
			Required: []string{"dataset_id", "title", "urls"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		datasetID, _ := args["dataset_id"].(string)
		title, _ := args["title"].(string)
		if apiKey == "" || datasetID == "" || title == "" {
			return nil, fmt.Errorf("api_key, dataset_id, and title are required")
		}
		var urls []string
		if arr, ok := args["urls"].([]interface{}); ok {
			for _, v := range arr {
				if s, ok := v.(string); ok && s != "" {
					urls = append(urls, s)
				}
			}
		}
		if len(urls) == 0 {
			return nil, fmt.Errorf("at least one URL is required")
		}
		return gw.ImportKnowledgeFromURL(ctx, apiKey, datasetID, title, urls)
	})
}

func registerImportKnowledgeWebsite(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "import_knowledge_from_website",
		Description: "Crawl an entire website via BFS and import all discovered pages into the knowledge base. Best for documentation sites, help centers, company websites. Requires a valid dataset_id UUID — call create_knowledge_base first.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"dataset_id": {Type: "string", Description: "Knowledge base (dataset) UUID from create_knowledge_base"},
				"url":        {Type: "string", Description: "Website root URL to crawl"},
				"max_pages":  {Type: "integer", Description: "Max pages to crawl (default: no limit)"},
			},
			Required: []string{"dataset_id", "url"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		datasetID, _ := args["dataset_id"].(string)
		webURL, _ := args["url"].(string)
		if apiKey == "" || datasetID == "" || webURL == "" {
			return nil, fmt.Errorf("api_key, dataset_id, and url are required")
		}
		maxPages := 0
		if d, ok := args["max_pages"].(float64); ok && d > 0 {
			maxPages = int(d)
		}
		return gw.ImportKnowledgeFromWebsite(ctx, apiKey, datasetID, webURL, maxPages)
	})
}

func registerGetImportStatus(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_import_status",
		Description: "Check the status of a knowledge base import job. Use this after importing from URL to see if indexing is complete. The 'url' parameter must match the original import URL.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"dataset_id": {Type: "string", Description: "Knowledge base (dataset) ID"},
				"job_id":     {Type: "string", Description: "Import job ID returned from import call"},
				"url":        {Type: "string", Description: "The original import URL (required by the RAG backend)"},
			},
			Required: []string{"dataset_id", "job_id", "url"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		datasetID, _ := args["dataset_id"].(string)
		jobID, _ := args["job_id"].(string)
		importURL, _ := args["url"].(string)
		if apiKey == "" || datasetID == "" || jobID == "" {
			return nil, fmt.Errorf("api_key, dataset_id, and job_id are required")
		}
		return gw.GetImportStatus(ctx, apiKey, datasetID, jobID, importURL)
	})
}

func registerCreateKnowledgeBase(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "create_knowledge_base",
		Description: "Create a new empty knowledge base. Returns the dataset_id (UUID) needed for importing documents via import_knowledge_from_url, import_knowledge_from_website, or listing documents. Always create the KB first, then import content into it.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"name":        {Type: "string", Description: "Knowledge base name"},
				"description": {Type: "string", Description: "Knowledge base description"},
			},
			Required: []string{"name"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		name, _ := args["name"].(string)
		if apiKey == "" || name == "" {
			return nil, fmt.Errorf("api_key and name are required")
		}
		req := gateway.KnowledgeBaseCreateRequest{
			Name:        name,
			Description: strval(args, "description"),
		}
		return gw.CreateKnowledgeBase(ctx, apiKey, req)
	})
}

func registerDeleteDocument(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "delete_knowledge_document",
		Description: "Delete a document from the knowledge base. This permanently removes the document and its indexed content. The 'url' field is the document's source URL from data_source_info — call list_knowledge_documents first to obtain it.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"dataset_id":  {Type: "string", Description: "Knowledge base (dataset) ID"},
				"document_id": {Type: "string", Description: "Document ID to delete"},
				"url":         {Type: "string", Description: "Source URL of the document (from data_source_info in list_knowledge_documents response)"},
			},
			Required: []string{"dataset_id", "document_id", "url"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		datasetID, _ := args["dataset_id"].(string)
		documentID, _ := args["document_id"].(string)
		docURL, _ := args["url"].(string)
		if apiKey == "" || datasetID == "" || documentID == "" || docURL == "" {
			return nil, fmt.Errorf("api_key, dataset_id, document_id, and url are required")
		}

		tenantID, err := gw.ResolveTenantID(ctx, apiKey)
		if err != nil {
			return nil, fmt.Errorf("resolve tenant_id for delete: %w", err)
		}

		body := &gateway.DocumentDeleteRequest{
			TenantID: tenantID,
			URLs:     []string{docURL},
			Size:     1,
		}
		if err := gw.DeleteDocument(ctx, apiKey, datasetID, documentID, body); err != nil {
			return nil, err
		}
		return map[string]string{"status": "deleted"}, nil
	})
}

func registerCreateDocument(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "create_knowledge_document",
		Description: "Create a text document in the knowledge base. The content will be chunked and indexed for retrieval.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"dataset_id": {Type: "string", Description: "Knowledge base (dataset) ID"},
				"title":      {Type: "string", Description: "Document title (max 255 chars)"},
				"content":    {Type: "string", Description: "Document text content"},
			},
			Required: []string{"dataset_id", "title", "content"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		datasetID, _ := args["dataset_id"].(string)
		title, _ := args["title"].(string)
		content, _ := args["content"].(string)
		if apiKey == "" || datasetID == "" || title == "" || content == "" {
			return nil, fmt.Errorf("api_key, dataset_id, title, and content are required")
		}
		req := gateway.DocumentCreateRequest{Title: title, Content: content}
		return gw.CreateDocument(ctx, apiKey, datasetID, req)
	})
}
