package openapi

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

// DocKind identifies which generated document an endpoint belongs to.
type DocKind string

const (
	// ControlPlaneAKSK is the control-plane document authenticated by
	// TiDB Cloud API keys (X-TiDBCloud-Public-Key / X-TiDBCloud-Private-Key).
	ControlPlaneAKSK DocKind = "control-plane-aksk"
	// ControlPlaneJWT is the control-plane document authenticated by a
	// TiDB Cloud user JWT (X-TiDB-Cloud-Authorization: Bearer <jwt>) with
	// the organization id passed as the orgId query parameter.
	ControlPlaneJWT DocKind = "control-plane-jwt"
	// DataPlane is the tenant document authenticated by a drive9 API key
	// (Authorization: Bearer <key> or X-API-Key).
	DataPlane DocKind = "data-plane"
)

// Security scheme names emitted into components.securitySchemes.
const (
	// SecurityTiDBCloudPublicKey / SecurityTiDBCloudPrivateKey are the two
	// halves of TiDB Cloud API-key authentication.
	SecurityTiDBCloudPublicKey  = "TiDBCloudPublicKey"
	SecurityTiDBCloudPrivateKey = "TiDBCloudPrivateKey"
	// SecurityTiDBCloudJWT authenticates with a TiDB Cloud user JWT in the
	// X-TiDB-Cloud-Authorization header.
	SecurityTiDBCloudJWT = "TiDBCloudJWT"
	// SecurityDrive9Key authenticates with a drive9 API key.
	SecurityDrive9Key = "Drive9Key"
	// SecurityVaultCapability authenticates with a vault capability or
	// grant token (vault_ / vt_ prefixes).
	SecurityVaultCapability = "VaultCapability"
)

// Endpoint describes one HTTP operation in the registry. The same Endpoint
// can be included in several documents via Docs.
type Endpoint struct {
	// Method is the HTTP verb (GET, POST, PUT, PATCH, DELETE, HEAD).
	Method string
	// Path is the path template, e.g. "/v1/fs/{path}".
	Path string
	// Summary is a short human-readable title.
	Summary string
	// Description is an optional longer explanation.
	Description string
	// OperationID is optional; generated from Method+Path when empty.
	OperationID string
	// Tags group operations in the document; empty means "default".
	Tags []string
	// Parameters are the explicit path/query/header parameters.
	Parameters []Parameter
	// RequestBody is the JSON request body schema; nil means no body.
	RequestBody *SchemaRef
	// requestBodyRequired marks the body as required. A nil value means
	// "required by default" (true); body builders may explicitly set it false.
	// *bool is used so the zero value of a directly-constructed Endpoint
	// still means "required", matching the builder helpers.
	requestBodyRequired *bool
	// RequestBodyDescription is rendered as the requestBody.description; the
	// builder helpers fill it from the WithBody argument.
	RequestBodyDescription string
	// RequestMediaType overrides the default "application/json" (e.g.
	// "application/octet-stream" for raw file bodies).
	RequestMediaType string
	// RequestContent declares multiple media types, taking precedence over
	// RequestBody/RequestMediaType when set.
	RequestContent []MediaContent
	// ClientErrorSchema overrides the default 4XX error schema for this operation.
	ClientErrorSchema *SchemaRef
	// Responses lists the success responses. A default error response
	// referencing ErrorResponse is appended automatically.
	Responses []Response
	// Security overrides the document-level security with the named schemes
	// (e.g. SecurityVaultCapability for /v1/vault/read). Empty means the
	// document-level security applies.
	Security []string
	// Docs lists the documents this endpoint appears in.
	Docs []DocKind
}

// Parameter describes a single path/query/header parameter.
type Parameter struct {
	// In is one of "path", "query", "header".
	In string
	// Name is the parameter name.
	Name string
	// Required marks mandatory parameters (always true for path).
	Required bool
	// Schema is the parameter value schema (scalar).
	Schema SchemaRef
	// Description is an optional human-readable explanation.
	Description string
}

// MediaContent is one media type entry of a request or response.
type MediaContent struct {
	// MediaType is the content type, e.g. "application/json" or
	// "application/octet-stream".
	MediaType string
	// Schema is the body schema for this media type; nil means no body.
	Schema *SchemaRef
}

// Response describes one HTTP response.
type Response struct {
	// Code is the HTTP status code (200, 201, 202, 204, ...).
	Code int
	// Description is the response meaning; defaults to the status text.
	Description string
	// Schema is the response body schema; nil means no body. When Content
	// is set, Schema/MediaType are ignored.
	Schema *SchemaRef
	// MediaType overrides the default "application/json"; empty with a
	// non-nil Schema means application/json. When Schema is nil and
	// MediaType is "application/octet-stream" a binary response is emitted.
	MediaType string
	// Content lists multiple media types for the same status code (e.g. a
	// binary body and a JSON body). When set, it takes precedence over
	// Schema/MediaType.
	Content []MediaContent
}

// Doc describes one generated OpenAPI document.
type Doc struct {
	// Kind identifies the document (controls output keying in the server
	// package). Empty is allowed for standalone builds.
	Kind DocKind
	// Title is the OpenAPI info.title.
	Title string
	// Description is the OpenAPI info.description.
	Description string
	// Version is the OpenAPI info.version (defaults to "1.0.0").
	Version string
	// Servers is the servers block; may be empty.
	Servers []Server
	// Security is the document-level security requirements.
	Security []map[string][]string
	// Endpoints lists all endpoints of this document.
	Endpoints []Endpoint
	// ExtraSchemas registers additional components (e.g. ErrorResponse).
	ExtraSchemas []SchemaRef
	// OrgIDQuery forces a required orgId query parameter onto every
	// endpoint (control-plane JWT document).
	OrgIDQuery bool
}

// Server is an OpenAPI server entry.
type Server struct {
	URL         string
	Description string
}

// ErrorResponse is the shared error body schema: {"error": string}.
var ErrorResponse = Raw("ErrorResponse", jsonSchema{
	"type":       "object",
	"properties": jsonSchema{"error": jsonSchema{"type": "string"}},
	"required":   []string{"error"},
})

// securitySchemes are the component definitions shared by all documents.
var securitySchemes = map[string]any{
	SecurityTiDBCloudPublicKey: map[string]any{
		"type":        "apiKey",
		"in":          "header",
		"name":        "X-TiDBCloud-Public-Key",
		"description": "TiDB Cloud API key public key (access key id).",
	},
	SecurityTiDBCloudPrivateKey: map[string]any{
		"type":        "apiKey",
		"in":          "header",
		"name":        "X-TiDBCloud-Private-Key",
		"description": "TiDB Cloud API key private key (secret key).",
	},
	SecurityTiDBCloudJWT: map[string]any{
		"type":        "apiKey",
		"in":          "header",
		"name":        "X-TiDB-Cloud-Authorization",
		"description": "TiDB Cloud user JWT: `Bearer <jwt>`. The target organization must be passed as the `orgId` query parameter on every request.",
	},
	SecurityDrive9Key: map[string]any{
		"type":         "http",
		"scheme":       "bearer",
		"bearerFormat": "drive9-api-key",
		"description":  "drive9 API key. Send as `Authorization: Bearer <key>` (the `X-API-Key` header is accepted as an alternative).",
	},
	SecurityVaultCapability: map[string]any{
		"type":         "http",
		"scheme":       "bearer",
		"bearerFormat": "vault-capability-token",
		"description":  "Vault capability token (`vault_...` prefix) or delegated grant token (`vt_...` prefix).",
	},
}

// Build renders a Doc into an OpenAPI 3.0.3 JSON document.
func Build(d Doc) ([]byte, error) {
	if d.Version == "" {
		d.Version = "1.0.0"
	}
	registry := newTypeRegistry()
	// Register the shared error schema first so it is always present.
	if _, err := registry.resolve(ErrorResponse); err != nil {
		return nil, err
	}
	for _, s := range d.ExtraSchemas {
		if _, err := registry.resolve(s); err != nil {
			return nil, err
		}
	}

	paths := make(map[string]any)
	usedSchemes := make(map[string]bool)
	for _, req := range d.Security {
		for scheme := range req {
			usedSchemes[scheme] = true
		}
	}
	for _, ep := range d.Endpoints {
		if err := validateEndpoint(ep); err != nil {
			return nil, err
		}
		for _, scheme := range ep.Security {
			usedSchemes[scheme] = true
		}
		op, err := buildOperation(ep, d.OrgIDQuery, registry)
		if err != nil {
			return nil, fmt.Errorf("endpoint %s %s: %w", ep.Method, ep.Path, err)
		}
		pathOps, ok := paths[ep.Path].(map[string]any)
		if !ok {
			pathOps = make(map[string]any)
			paths[ep.Path] = pathOps
		}
		key := strings.ToLower(ep.Method)
		if _, exists := pathOps[key]; exists {
			return nil, fmt.Errorf("duplicate operation %s %s", ep.Method, ep.Path)
		}
		pathOps[key] = op
	}

	components := map[string]any{
		"schemas": registry.schemas,
	}
	usedSchemesMap := make(map[string]any, len(usedSchemes))
	for scheme := range usedSchemes {
		def, ok := securitySchemes[scheme]
		if !ok {
			return nil, fmt.Errorf("unknown security scheme %q", scheme)
		}
		usedSchemesMap[scheme] = def
	}
	if len(usedSchemesMap) > 0 {
		components["securitySchemes"] = usedSchemesMap
	}
	if len(components["schemas"].(map[string]jsonSchema)) == 0 {
		delete(components, "schemas")
	}
	if err := registry.checkRefs(); err != nil {
		return nil, err
	}

	doc := map[string]any{
		"openapi": "3.0.3",
		"info": map[string]any{
			"title":       d.Title,
			"description": d.Description,
			"version":     d.Version,
			"license": map[string]any{
				"name": "Apache-2.0",
				"url":  "https://www.apache.org/licenses/LICENSE-2.0.html",
			},
		},
		"paths":      paths,
		"components": components,
	}
	if len(d.Servers) > 0 {
		servers := make([]map[string]any, 0, len(d.Servers))
		for _, s := range d.Servers {
			entry := map[string]any{"url": s.URL}
			if s.Description != "" {
				entry["description"] = s.Description
			}
			servers = append(servers, entry)
		}
		doc["servers"] = servers
	}
	if len(d.Security) > 0 {
		doc["security"] = d.Security
	}
	// Validate every $ref in the assembled document, including unnamed Raw
	// schemas inlined into paths that never become components.
	if err := checkDocumentRefs("document", doc, registry.schemas); err != nil {
		return nil, err
	}
	return json.MarshalIndent(doc, "", "  ")
}

func validateEndpoint(ep Endpoint) error {
	if ep.Method == "" {
		return fmt.Errorf("endpoint %s: missing method", ep.Path)
	}
	if ep.Path == "" {
		return fmt.Errorf("endpoint %s: missing path", ep.Method)
	}
	for _, p := range ep.Parameters {
		switch p.In {
		case "path", "query", "header":
		default:
			return fmt.Errorf("endpoint %s %s: invalid parameter location %q", ep.Method, ep.Path, p.In)
		}
		if p.In == "path" {
			if !strings.Contains(ep.Path, "{"+p.Name+"}") {
				return fmt.Errorf("endpoint %s %s: path parameter %q not present in path template", ep.Method, ep.Path, p.Name)
			}
			if !p.Required {
				return fmt.Errorf("endpoint %s %s: path parameter %q must be required", ep.Method, ep.Path, p.Name)
			}
		}
	}
	return nil
}

func buildOperation(ep Endpoint, orgIDQuery bool, registry *typeRegistry) (map[string]any, error) {
	op := map[string]any{
		"summary": ep.Summary,
	}
	if ep.Description != "" {
		op["description"] = ep.Description
	}
	operationID := ep.OperationID
	if operationID == "" {
		operationID = operationIDFor(ep.Method, ep.Path)
	}
	op["operationId"] = operationID
	if len(ep.Tags) > 0 {
		op["tags"] = ep.Tags
	}
	if len(ep.Security) > 0 {
		security := make([]map[string]any, 0, len(ep.Security))
		for _, scheme := range ep.Security {
			security = append(security, map[string]any{scheme: []any{}})
		}
		op["security"] = security
	}

	parameters := make([]map[string]any, 0, len(ep.Parameters)+1)
	for _, p := range ep.Parameters {
		pSchema, err := registry.resolve(p.Schema)
		if err != nil {
			return nil, fmt.Errorf("parameter %s %s: %w", p.Name, p.In, err)
		}
		param := map[string]any{
			"name":     p.Name,
			"in":       p.In,
			"required": p.Required,
			"schema":   pSchema,
		}
		if p.Description != "" {
			param["description"] = p.Description
		}
		parameters = append(parameters, param)
	}
	if orgIDQuery {
		parameters = append(parameters, map[string]any{
			"name":        "orgId",
			"in":          "query",
			"required":    true,
			"description": "TiDB Cloud organization id the caller acts on. Must be a numeric non-zero value.",
			"schema":      map[string]any{"type": "integer", "format": "int64"},
		})
	}
	if len(parameters) > 0 {
		sort.Slice(parameters, func(i, j int) bool {
			a, b := parameters[i], parameters[j]
			if a["in"] != b["in"] {
				return a["in"].(string) < b["in"].(string)
			}
			return a["name"].(string) < b["name"].(string)
		})
		op["parameters"] = parameters
	}

	if ep.RequestBody != nil || len(ep.RequestContent) > 0 {
		required := true
		if ep.requestBodyRequired != nil {
			required = *ep.requestBodyRequired
		}
		contents := ep.RequestContent
		if len(contents) == 0 {
			mediaType := ep.RequestMediaType
			if mediaType == "" {
				mediaType = "application/json"
			}
			if mediaType != "application/json" {
				// Preserve the existing single-media binary body requirement.
				required = true
			}
			contents = []MediaContent{{MediaType: mediaType, Schema: ep.RequestBody}}
		}
		content := map[string]any{}
		for _, c := range contents {
			if c.MediaType == "" || c.Schema == nil {
				return nil, fmt.Errorf("request body requires a media type and schema")
			}
			if _, exists := content[c.MediaType]; exists {
				return nil, fmt.Errorf("duplicate request media type %q", c.MediaType)
			}
			if c.MediaType == "application/octet-stream" {
				content[c.MediaType] = map[string]any{"schema": map[string]any{"type": "string", "format": "binary"}}
				continue
			}
			bodySchema, err := registry.resolve(*c.Schema)
			if err != nil {
				return nil, fmt.Errorf("request body: %w", err)
			}
			content[c.MediaType] = map[string]any{"schema": bodySchema}
		}
		body := map[string]any{
			"required": required,
			"content":  content,
		}
		if ep.RequestBodyDescription != "" {
			body["description"] = ep.RequestBodyDescription
		}
		op["requestBody"] = body
	}

	responses := make(map[string]any)
	for _, resp := range ep.Responses {
		key := fmt.Sprintf("%d", resp.Code)
		if _, exists := responses[key]; exists {
			return nil, fmt.Errorf("duplicate response status %d", resp.Code)
		}
		entry, err := buildResponse(resp, registry)
		if err != nil {
			return nil, err
		}
		responses[key] = entry
	}
	// Range responses keep the shared error shape explicit so code
	// generators emit a usable error type.
	errorSchema, err := registry.resolve(ErrorResponse)
	if err != nil {
		return nil, err
	}
	errorEntry := map[string]any{
		"description": "Error. Body: {\"error\": string}.",
		"content": map[string]any{
			"application/json": map[string]any{"schema": errorSchema},
		},
	}
	if _, ok := responses["4XX"]; !ok {
		responses["4XX"] = errorEntry
		if ep.ClientErrorSchema != nil {
			entry, err := buildResponse(Response{Schema: ep.ClientErrorSchema, Description: "Client error. See schema for error fields."}, registry)
			if err != nil {
				return nil, err
			}
			responses["4XX"] = entry
		}
	}
	if _, ok := responses["5XX"]; !ok {
		responses["5XX"] = errorEntry
	}
	op["responses"] = responses
	return op, nil
}

func buildResponse(resp Response, registry *typeRegistry) (map[string]any, error) {
	desc := resp.Description
	if desc == "" {
		desc = fmt.Sprintf("%d %s", resp.Code, statusText(resp.Code))
	}
	entry := map[string]any{"description": desc}

	contents := resp.Content
	if len(contents) == 0 {
		if resp.Schema == nil {
			// Bodyless response (204, 302, SSE stream, header-only stat):
			// no content block unless the registry explicitly declared one.
			if resp.MediaType == "" {
				return entry, nil
			}
			contents = []MediaContent{{MediaType: resp.MediaType}}
		} else {
			mediaType := resp.MediaType
			if mediaType == "" {
				mediaType = "application/json"
			}
			contents = []MediaContent{{MediaType: mediaType, Schema: resp.Schema}}
		}
	}
	content := make(map[string]any, len(contents))
	for _, c := range contents {
		if c.MediaType == "" {
			return nil, fmt.Errorf("response %d: empty media type", resp.Code)
		}
		if c.Schema == nil {
			// Bodyless media type declaration (e.g. text/event-stream).
			content[c.MediaType] = map[string]any{}
			continue
		}
		if c.MediaType == "application/octet-stream" {
			content[c.MediaType] = map[string]any{"schema": map[string]any{"type": "string", "format": "binary"}}
			continue
		}
		schema, err := registry.resolve(*c.Schema)
		if err != nil {
			return nil, fmt.Errorf("response %d content %s: %w", resp.Code, c.MediaType, err)
		}
		content[c.MediaType] = map[string]any{"schema": schema}
	}
	if len(content) > 0 {
		entry["content"] = content
	}
	return entry, nil
}

func statusText(code int) string {
	switch code {
	case 200:
		return "OK"
	case 201:
		return "Created"
	case 202:
		return "Accepted"
	case 204:
		return "No Content"
	default:
		return "Response"
	}
}

func operationIDFor(method, path string) string {
	var b strings.Builder
	b.WriteString(strings.ToLower(method))
	for _, segment := range strings.Split(strings.Trim(path, "/"), "/") {
		if segment == "" {
			continue
		}
		segment = strings.ReplaceAll(segment, ":", "-")
		if strings.HasPrefix(segment, "{") && strings.HasSuffix(segment, "}") {
			segment = "by-" + strings.TrimSuffix(strings.TrimPrefix(segment, "{"), "}")
		}
		b.WriteString("-")
		b.WriteString(segment)
	}
	return b.String()
}

// --- Builder DSL helpers (used by the registry in pkg/server) ---------------

// WithParams appends parameters to the endpoint.
func (e Endpoint) WithParams(params ...Parameter) Endpoint {
	e.Parameters = append(e.Parameters, params...)
	return e
}

// WithDescription sets the endpoint description.
func (e Endpoint) WithDescription(desc string) Endpoint {
	e.Description = desc
	return e
}

// WithBody sets the JSON request body. mediaTypes may override the default
// "application/json" (e.g. "application/octet-stream").
func (e Endpoint) WithBody(schema *SchemaRef, desc string, mediaTypes ...string) Endpoint {
	e.RequestBody = schema
	e.RequestBodyDescription = desc
	if len(mediaTypes) > 0 {
		e.RequestMediaType = mediaTypes[0]
	}
	return e
}

// WithOptionalBody sets a non-required request body.
func (e Endpoint) WithOptionalBody(schema *SchemaRef, desc string, mediaTypes ...string) Endpoint {
	e = e.WithBody(schema, desc, mediaTypes...)
	required := false
	e.requestBodyRequired = &required
	return e
}

// WithBodyContent declares alternative request media types and their schemas.
func (e Endpoint) WithBodyContent(desc string, required bool, content ...MediaContent) Endpoint {
	e.RequestContent = content
	e.RequestBodyDescription = desc
	e.requestBodyRequired = &required
	return e
}

// WithClientError sets this operation's 4XX schema, leaving 5XX errors unchanged.
func (e Endpoint) WithClientError(schema *SchemaRef) Endpoint {
	e.ClientErrorSchema = schema
	return e
}

// WithResponses sets the success responses.
func (e Endpoint) WithResponses(responses []Response) Endpoint {
	e.Responses = append(e.Responses, responses...)
	return e
}

// WithHeaders adds optional header parameters.
func (e Endpoint) WithHeaders(names ...string) Endpoint {
	for _, name := range names {
		e.Parameters = append(e.Parameters, Parameter{
			In:          "header",
			Name:        name,
			Description: "Optional request header.",
			Schema:      String(),
		})
	}
	return e
}

// WithRequiredHeaders adds required header parameters.
func (e Endpoint) WithRequiredHeaders(names ...string) Endpoint {
	for _, name := range names {
		e.Parameters = append(e.Parameters, Parameter{
			In:          "header",
			Name:        name,
			Required:    true,
			Description: "Required request header.",
			Schema:      String(),
		})
	}
	return e
}

// WithSecurity overrides the document-level security for this endpoint.
func (e Endpoint) WithSecurity(schemes ...string) Endpoint {
	e.Security = append(e.Security, schemes...)
	return e
}