package server

import (
	"fmt"
	"reflect"

	"github.com/mem9-ai/drive9/pkg/backend"
	"github.com/mem9-ai/drive9/pkg/datastore"
	"github.com/mem9-ai/drive9/pkg/openapi"
	"github.com/mem9-ai/drive9/pkg/vault"
)

// This file assembles the three generated OpenAPI documents from the shared
// endpoint registry (openapiRegistry in openapi_registry.go).

// OpenAPIDocuments builds the three OpenAPI 3.0.3 documents for the drive9
// API: control plane with TiDB Cloud API keys, control plane with TiDB Cloud
// user JWTs (orgId query parameter), and the data plane with drive9 API
// keys. The returned map is keyed by document kind.
//
// serverURL is the externally reachable base URL advertised in the
// documents' servers block; an empty string leaves the servers block out.
func OpenAPIDocuments(serverURL string) (map[openapi.DocKind][]byte, error) {
	registry := openapiRegistry()

	var servers []openapi.Server
	if serverURL != "" {
		servers = []openapi.Server{{URL: serverURL}}
	}

	controlPlaneDescription := "Control-plane API for managing drive9 tenants, quotas, " +
		"object backends, tenant pools and API keys from TiDB Cloud."

	docs := []openapi.Doc{
		{
			Kind:  openapi.ControlPlaneAKSK,
			Title: "Drive9 Control Plane API — TiDB Cloud API Key",
			Description: controlPlaneDescription + " Authenticated with a TiDB Cloud API key pair " +
				"(X-TiDBCloud-Public-Key / X-TiDBCloud-Private-Key headers).",
			Version: "1.0.0",
			Servers: servers,
			Security: []map[string][]string{
				{openapi.SecurityTiDBCloudPublicKey: {}, openapi.SecurityTiDBCloudPrivateKey: {}},
			},
			Endpoints: filterEndpoints(registry, openapi.ControlPlaneAKSK),
		},
		{
			Kind:  openapi.ControlPlaneJWT,
			Title: "Drive9 Control Plane API — TiDB Cloud JWT",
			Description: controlPlaneDescription + " Authenticated with a TiDB Cloud user JWT " +
				"(X-TiDB-Cloud-Authorization: Bearer <jwt>). Every request must carry the numeric " +
				"orgId query parameter selecting the organization.",
			Version: "1.0.0",
			Servers: servers,
			Security: []map[string][]string{
				{openapi.SecurityTiDBCloudJWT: {}},
			},
			Endpoints:  filterEndpoints(registry, openapi.ControlPlaneJWT),
			OrgIDQuery: true,
		},
		{
			Kind:  openapi.DataPlane,
			Title: "Drive9 Data Plane API",
			Description: "Filesystem, upload, journal, vault, git workspace and layer APIs. " +
				"Authenticated with a drive9 API key (Authorization: Bearer <key>; the X-API-Key " +
				"header is accepted as an alternative). The /v1/vault/read endpoints instead use a " +
				"vault capability token or delegated grant token.",
			Version: "1.0.0",
			Servers: servers,
			Security: []map[string][]string{
				{openapi.SecurityDrive9Key: {}},
			},
			Endpoints: filterEndpoints(registry, openapi.DataPlane),
			// Types referenced by hand-written schemas but never used as a
			// direct endpoint body/response; register so $refs resolve.
			ExtraSchemas: []openapi.SchemaRef{
				openapi.Ref(reflect.TypeOf(backend.PresignChecksum{})),
				openapi.Ref(reflect.TypeOf(backend.PresignPartEntry{})),
				openapi.Ref(reflect.TypeOf(fsLayerEventResponse{})),
				openapi.Ref(reflect.TypeOf(vault.AuditEvent{})),
				openapi.Ref(reflect.TypeOf(datastore.SearchResult{})),
				openapi.Ref(reflect.TypeOf(fsListResponseEntry{})),
				openapi.Ref(reflect.TypeOf(fsStatMetadataResponse{})),
			},
		},
	}

	out := make(map[openapi.DocKind][]byte, len(docs))
	for _, doc := range docs {
		raw, err := openapi.Build(doc)
		if err != nil {
			return nil, fmt.Errorf("build %s document: %w", doc.Kind, err)
		}
		out[doc.Kind] = raw
	}
	return out, nil
}

func filterEndpoints(registry []openapi.Endpoint, kind openapi.DocKind) []openapi.Endpoint {
	var out []openapi.Endpoint
	for _, ep := range registry {
		for _, k := range ep.Docs {
			if k == kind {
				out = append(out, ep)
				break
			}
		}
	}
	return out
}