package dify

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
)

// CreateDataset creates a new knowledge base (dataset) in Dify.
func (c *Client) CreateDataset(ctx context.Context, apiKey string, req *DatasetCreateRequest) (*Dataset, error) {
	var result Dataset
	if err := c.doJSON(ctx, "POST", "/datasets", apiKey, req, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

// DeleteDataset deletes a knowledge base.
func (c *Client) DeleteDataset(ctx context.Context, apiKey, datasetID string) error {
	return c.doJSON(ctx, "DELETE", "/datasets/"+datasetID, apiKey, nil, nil)
}

// UploadDocument uploads a file to a dataset for indexing.
func (c *Client) UploadDocument(ctx context.Context, apiKey, datasetID string, filename string, fileReader io.Reader) (*DocumentIndexingStatus, error) {
	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)

	part, err := writer.CreateFormFile("file", filename)
	if err != nil {
		return nil, fmt.Errorf("create form file: %w", err)
	}
	if _, err := io.Copy(part, fileReader); err != nil {
		return nil, fmt.Errorf("copy file: %w", err)
	}

	// indexing mode
	if err := writer.WriteField("indexing_technique", "high_quality"); err != nil {
		return nil, err
	}
	if err := writer.WriteField("process_rule", `{"mode": "automatic"}`); err != nil {
		return nil, err
	}
	writer.Close()

	path := fmt.Sprintf("/datasets/%s/document/create_by_file", datasetID)
	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+path, &buf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", writer.FormDataContentType())

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("upload failed (status %d): %s", resp.StatusCode, string(body))
	}

	// Dify returns document info in response
	var result struct {
		Document DocumentIndexingStatus `json:"document"`
	}
	if err := decodeJSON(resp.Body, &result); err != nil {
		return nil, err
	}
	return &result.Document, nil
}

// DeleteDocument deletes a document from a dataset.
func (c *Client) DeleteDocument(ctx context.Context, apiKey, datasetID, documentID string) error {
	path := fmt.Sprintf("/datasets/%s/documents/%s", datasetID, documentID)
	return c.doJSON(ctx, "DELETE", path, apiKey, nil, nil)
}

// GetDocumentIndexingStatus checks the indexing status of a document.
func (c *Client) GetDocumentIndexingStatus(ctx context.Context, apiKey, datasetID, batch string) (*DocumentIndexingStatus, error) {
	path := fmt.Sprintf("/datasets/%s/documents/%s/indexing-status", datasetID, batch)
	var result DocumentIndexingStatus
	if err := c.doJSON(ctx, "GET", path, apiKey, nil, &result); err != nil {
		return nil, err
	}
	return &result, nil
}

func decodeJSON(r io.Reader, v any) error {
	return json.NewDecoder(r).Decode(v)
}