package knowledgeqa

import (
	"context"
	"errors"
	"net/http"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/api/meshxapi"
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
	"github.com/tidbcloud/AgentTeam/internal/service/knowledgeqa/runtime"
)

const turnExecutionTimeout = runtime.DefaultTurnExecutionTimeout

type ExecuteRequest struct {
	OrganizationID   string                 `json:"organizationId"`
	TurnID           string                 `json:"turnId"`
	Question         string                 `json:"question"`
	KnowledgeBaseIDs []string               `json:"knowledgeBaseIds,omitempty"`
	SourceIDs        []string               `json:"sourceIds,omitempty"`
	RootHandleIDs    []string               `json:"rootHandleIds,omitempty"`
	WebSearchEnabled bool                   `json:"webSearchEnabled"`
	Model            string                 `json:"model"`
	ReasonLevel      string                 `json:"reasonLevel"`
	MaaSCapability   string                 `json:"maasCapability"`
	IdempotencyKey   string                 `json:"idempotencyKey"`
	RuntimeContext   runtime.RuntimeContext `json:"runtimeContext"`
}

type PrepareRequest struct {
	OrganizationID   string   `json:"organizationId"`
	TurnID           string   `json:"turnId"`
	KnowledgeBaseIDs []string `json:"knowledgeBaseIds,omitempty"`
	SourceIDs        []string `json:"sourceIds,omitempty"`
}

type OpenCitationRequest struct {
	OrganizationID          string                   `json:"organizationId"`
	TurnID                  string                   `json:"turnId"`
	Citation                runtime.CitationMetadata `json:"citation"`
	AllowRestrictedEvidence bool                     `json:"allowRestrictedEvidence"`
}

// RecoverAccountingRequest identifies one exact attempt under the authenticated
// internal requester and workspace authority.
type RecoverAccountingRequest struct {
	OrganizationID string `json:"organizationId"`
	TurnID         string `json:"turnId"`
	AttemptID      string `json:"attemptId"`
}

type Handler struct {
	Runtime *runtime.Runtime
}

func RegisterRuntimeRoutes(router *http.ServeMux, handler *Handler, middleware func(http.Handler) http.Handler) {
	if router == nil || handler == nil {
		return
	}
	route := meshxapi.HandleAuthJSONLimit[ExecuteRequest](1<<20, handler.Execute, errorHandling())
	if middleware != nil {
		route = middleware(route)
	}
	router.Handle("POST /api/v1/kb/knowledge-qa/runtime/execute", route)
	recoverAccounting := meshxapi.HandleAuthJSONLimit[RecoverAccountingRequest](16<<10, handler.RecoverAccounting, errorHandling())
	if middleware != nil {
		recoverAccounting = middleware(recoverAccounting)
	}
	router.Handle("POST /api/v1/kb/knowledge-qa/runtime/accounting/recover", recoverAccounting)
	prepare := meshxapi.HandleAuthJSONLimit[PrepareRequest](128<<10, handler.Prepare, errorHandling())
	if middleware != nil {
		prepare = middleware(prepare)
	}
	router.Handle("POST /api/v1/kb/knowledge-qa/runtime/prepare", prepare)
	openCitation := meshxapi.HandleAuthJSONLimit[OpenCitationRequest](64<<10, handler.OpenCitation, errorHandling())
	if middleware != nil {
		openCitation = middleware(openCitation)
	}
	router.Handle("POST /api/v1/kb/knowledge-qa/citations/open", openCitation)
}

// RecoverAccounting handles POST /api/v1/kb/knowledge-qa/runtime/accounting/recover.
// It returns only the owner-bound aggregate usage checkpoint for one exact
// attempt and never exposes model or repository content.
func (handler *Handler) RecoverAccounting(request meshxapi.AuthenticatedJSON[RecoverAccountingRequest]) (meshxapi.Response, error) {
	if handler == nil || handler.Runtime == nil {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusServiceUnavailable, "Knowledge Q&A runtime is unavailable")
	}
	if !hasInternalPurpose(request.Identity.RawClaims, "knowledge_qa_runtime_accounting") {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusForbidden, "Knowledge Q&A accounting recovery is server-internal")
	}
	workspaceID := strings.TrimSpace(request.WorkspaceID())
	requesterID := strings.TrimSpace(request.Subject())
	if workspaceID == "" || requesterID == "" {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusUnauthorized, "Knowledge Q&A requester authority is incomplete")
	}
	result, err := handler.Runtime.RecoverAccounting(request.Context(), runtime.RecoverAccountingRequest{
		Principal: protocol.Principal{
			OrganizationID: strings.TrimSpace(request.Body.OrganizationID),
			WorkspaceID:    workspaceID, RequesterID: requesterID,
			TurnID: strings.TrimSpace(request.Body.TurnID),
		},
		AttemptID: strings.TrimSpace(request.Body.AttemptID),
	})
	if err != nil {
		return meshxapi.Response{}, err
	}
	return meshxapi.OK(result), nil
}

func (handler *Handler) Prepare(request meshxapi.AuthenticatedJSON[PrepareRequest]) (meshxapi.Response, error) {
	if handler == nil || handler.Runtime == nil {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusServiceUnavailable, "Knowledge Q&A runtime is unavailable")
	}
	if !hasInternalPurpose(request.Identity.RawClaims, "knowledge_qa_runtime_prepare") {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusForbidden, "Knowledge Q&A runtime preparation is server-internal")
	}
	workspaceID := strings.TrimSpace(request.WorkspaceID())
	requesterID := strings.TrimSpace(request.Subject())
	if workspaceID == "" || requesterID == "" {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusUnauthorized, "Knowledge Q&A requester authority is incomplete")
	}
	result, err := handler.Runtime.Prepare(request.Context(), runtime.PrepareRequest{
		Principal: protocol.Principal{
			OrganizationID: strings.TrimSpace(request.Body.OrganizationID),
			WorkspaceID:    workspaceID,
			RequesterID:    requesterID,
			TurnID:         strings.TrimSpace(request.Body.TurnID),
		},
		KnowledgeBaseIDs: request.Body.KnowledgeBaseIDs,
		SourceIDs:        request.Body.SourceIDs,
	})
	if err != nil {
		return meshxapi.Response{}, err
	}
	return meshxapi.OK(result), nil
}

func (handler *Handler) Execute(request meshxapi.AuthenticatedJSON[ExecuteRequest]) (meshxapi.Response, error) {
	if handler == nil || handler.Runtime == nil {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusServiceUnavailable, "Knowledge Q&A runtime is unavailable")
	}
	if !hasInternalPurpose(request.Identity.RawClaims, "knowledge_qa_runtime") {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusForbidden, "Knowledge Q&A runtime execution is server-internal")
	}
	workspaceID := strings.TrimSpace(request.WorkspaceID())
	requesterID := strings.TrimSpace(request.Subject())
	if workspaceID == "" || requesterID == "" {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusUnauthorized, "Knowledge Q&A requester authority is incomplete")
	}
	ctx, cancel := contextWithRuntimeDeadline(request.Context())
	defer cancel()
	result, err := handler.Runtime.Execute(ctx, runtime.ExecuteRequest{
		Principal: protocol.Principal{
			OrganizationID: strings.TrimSpace(request.Body.OrganizationID),
			WorkspaceID:    workspaceID,
			RequesterID:    requesterID,
			TurnID:         strings.TrimSpace(request.Body.TurnID),
		},
		KnowledgeBaseIDs: request.Body.KnowledgeBaseIDs,
		SourceIDs:        request.Body.SourceIDs,
		RootHandleIDs:    request.Body.RootHandleIDs,
		WebSearchEnabled: request.Body.WebSearchEnabled,
		Question:         request.Body.Question,
		Model:            request.Body.Model,
		ReasonLevel:      request.Body.ReasonLevel,
		Capability:       request.Body.MaaSCapability,
		IdempotencyKey:   request.Body.IdempotencyKey,
		RuntimeContext:   request.Body.RuntimeContext,
	})
	if err != nil {
		return meshxapi.Response{}, err
	}
	return meshxapi.OK(result), nil
}

func (handler *Handler) OpenCitation(request meshxapi.AuthenticatedJSON[OpenCitationRequest]) (meshxapi.Response, error) {
	if handler == nil || handler.Runtime == nil {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusServiceUnavailable, "Knowledge Q&A runtime is unavailable")
	}
	if !hasInternalPurpose(request.Identity.RawClaims, "knowledge_qa_citation_open") {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusForbidden, "Knowledge Q&A citation opening is server-internal")
	}
	workspaceID := strings.TrimSpace(request.WorkspaceID())
	requesterID := strings.TrimSpace(request.Subject())
	if workspaceID == "" || requesterID == "" {
		return meshxapi.Response{}, meshxapi.NewHTTPError(http.StatusUnauthorized, "Knowledge Q&A requester authority is incomplete")
	}
	result, err := handler.Runtime.OpenCitation(request.Context(), runtime.OpenCitationRequest{
		Principal: protocol.Principal{
			OrganizationID: strings.TrimSpace(request.Body.OrganizationID),
			WorkspaceID:    workspaceID, RequesterID: requesterID,
			TurnID: strings.TrimSpace(request.Body.TurnID),
		},
		Citation:                request.Body.Citation,
		AllowRestrictedEvidence: request.Body.AllowRestrictedEvidence,
	})
	if err != nil {
		return meshxapi.Response{}, err
	}
	return meshxapi.OK(result), nil
}

func hasInternalPurpose(claims map[string]any, expected string) bool {
	value, ok := claims["agentteam_internal_purpose"].(string)
	return ok && strings.TrimSpace(value) == expected
}

func contextWithRuntimeDeadline(ctx context.Context) (context.Context, context.CancelFunc) {
	if _, ok := ctx.Deadline(); ok {
		return context.WithCancel(ctx)
	}
	return context.WithTimeout(ctx, turnExecutionTimeout)
}

func errorHandling() meshxapi.HandlerOption {
	return meshxapi.WithErrorClassifier(func(err error) (meshxapi.ErrorResponse, bool) {
		var runtimeErr *runtime.Error
		if !errors.As(err, &runtimeErr) {
			return meshxapi.ErrorResponse{}, false
		}
		status := http.StatusUnprocessableEntity
		switch runtimeErr.Code {
		case runtime.CodeInvalidRequest:
			status = http.StatusBadRequest
		case runtime.CodeHandleRejected:
			status = http.StatusForbidden
		case runtime.CodeCanceled:
			status = http.StatusRequestTimeout
		case runtime.CodeResourceExhausted:
			status = http.StatusTooManyRequests
		case runtime.CodeRepositoryFailed:
			status = http.StatusServiceUnavailable
		}
		return meshxapi.ErrorResponse{Status: status, Code: runtimeErr.Code, Message: runtimeErr.Message}, true
	})
}