package availability

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const maxResponseBytes = 1 << 20

// CreditAtomsPerCredit is the fixed MaaS conversion from exact credit atoms to
// the human-facing AgentTeam credit unit.
const CreditAtomsPerCredit = int64(1_000_000)

const (
	InventoryStatusActive     = "active_inventory"
	EntitlementStatusEligible = "eligible_entitled"
	BalanceStatusPositive     = "positive_balance"
	RuntimeStatusSupported    = "runtime_supported"

	ChatCompletionsProfileOpenAIRequiredToolsV1 = "chat_completions_openai_required_tools_v1"
	ChatCompletionsProfileDeepSeekV4ThinkingV1  = "chat_completions_deepseek_v4_thinking_v1"
)

// Client loads AgentTeam-facing hosted model availability from agentteam-maas internal APIs.
type Client struct {
	baseURL    string
	token      string
	httpClient *http.Client
}

// ClientError preserves the upstream HTTP status and bounded response body.
type ClientError struct {
	StatusCode int
	Body       string
}

func (err ClientError) Error() string {
	body := strings.TrimSpace(err.Body)
	return fmt.Sprintf("agentteam-maas hosted availability API returned HTTP %d: %s", err.StatusCode, body)
}

// NewClient constructs a agentteam-maas hosted availability client.
func NewClient(baseURL string, token string, httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return &Client{
		baseURL:    strings.TrimRight(strings.TrimSpace(baseURL), "/"),
		token:      strings.TrimSpace(token),
		httpClient: httpClient,
	}
}

// LoadHostedAvailability returns one organization's hosted inventory and status facts.
func (client *Client) LoadHostedAvailability(ctx context.Context, organizationID string) (Availability, error) {
	query := url.Values{}
	if trimmed := strings.TrimSpace(organizationID); trimmed != "" {
		query.Set("organization_id", trimmed)
	}
	var response Availability
	err := client.do(ctx, http.MethodGet, "/internal/maas/v1/hosted-availability", query, nil, &response)
	return response, err
}

// LoadHostedCatalog returns deployment-level inventory and executable
// capability facts without projecting any organization's entitlement or
// credits.
func (client *Client) LoadHostedCatalog(ctx context.Context) (Catalog, error) {
	var response Catalog
	err := client.do(ctx, http.MethodGet, "/internal/maas/v1/hosted-catalog", nil, nil, &response)
	return response, err
}

// UpsertHostedEntitlementGrant idempotently applies one AgentTeam-owned hosted entitlement grant.
func (client *Client) UpsertHostedEntitlementGrant(ctx context.Context, grantID string, payload EntitlementGrantRequest) (EntitlementGrant, error) {
	var response EntitlementGrant
	err := client.do(ctx, http.MethodPut, "/internal/maas/v1/hosted-entitlements/grants/"+url.PathEscape(strings.TrimSpace(grantID)), nil, payload, &response)
	return response, err
}

// HostedEntitlementGrantID returns the stable grant ID used by AgentTeam-created hosted grants.
func HostedEntitlementGrantID(organizationID string, targetKind string, targetID string) string {
	sum := sha256.Sum256([]byte(strings.TrimSpace(organizationID) + "\x00" + strings.TrimSpace(targetKind) + "\x00" + strings.TrimSpace(targetID)))
	return "ohg_agentteam_auto_" + hex.EncodeToString(sum[:16])
}

func (client *Client) do(ctx context.Context, method string, path string, query url.Values, body any, out any) error {
	if client == nil {
		return fmt.Errorf("agentteam-maas hosted availability client is not configured")
	}
	if strings.TrimSpace(client.baseURL) == "" {
		return fmt.Errorf("AGENTTEAM_MAAS_BASE_URL is required")
	}
	if strings.TrimSpace(client.token) == "" {
		return fmt.Errorf("AGENTTEAM_MAAS_INTERNAL_BEARER_TOKEN is required")
	}
	endpoint, err := client.endpoint(path, query)
	if err != nil {
		return err
	}
	var requestBody io.Reader
	if body != nil {
		encoded, err := json.Marshal(body)
		if err != nil {
			return fmt.Errorf("encode agentteam-maas hosted availability request: %w", err)
		}
		requestBody = bytes.NewReader(encoded)
	}
	request, err := http.NewRequestWithContext(ctx, method, endpoint, requestBody)
	if err != nil {
		return fmt.Errorf("build agentteam-maas hosted availability request: %w", err)
	}
	request.Header.Set("Authorization", "Bearer "+client.token)
	request.Header.Set("Accept", "application/json")
	if body != nil {
		request.Header.Set("Content-Type", "application/json")
	}

	httpClient := client.httpClient
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	response, err := httpClient.Do(request)
	if err != nil {
		return fmt.Errorf("call agentteam-maas hosted availability API: %w", err)
	}
	defer response.Body.Close()

	responseBody, err := io.ReadAll(io.LimitReader(response.Body, maxResponseBytes))
	if err != nil {
		return fmt.Errorf("read agentteam-maas hosted availability response: %w", err)
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return ClientError{StatusCode: response.StatusCode, Body: string(responseBody)}
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(responseBody, out); err != nil {
		return fmt.Errorf("decode agentteam-maas hosted availability response: %w", err)
	}
	return nil
}

func (client *Client) endpoint(path string, query url.Values) (string, error) {
	base, err := url.Parse(strings.TrimSpace(client.baseURL))
	if err != nil {
		return "", fmt.Errorf("parse AGENTTEAM_MAAS_BASE_URL: %w", err)
	}
	if base.Scheme == "" || base.Host == "" {
		return "", fmt.Errorf("AGENTTEAM_MAAS_BASE_URL must be an absolute URL")
	}
	base.Path = strings.TrimRight(base.Path, "/") + "/" + strings.TrimLeft(path, "/")
	base.RawQuery = query.Encode()
	return base.String(), nil
}

// Availability is the AgentTeam-facing hosted availability payload from agentteam-maas.
type Availability struct {
	Object               string    `json:"object"`
	OrganizationID       string    `json:"organizationId"`
	EntitlementStatus    string    `json:"entitlementStatus"`
	BalanceStatus        string    `json:"balanceStatus"`
	AvailableCreditAtoms int64     `json:"availableCreditAtoms"`
	CreditUnit           string    `json:"creditUnit"`
	UpdatedAt            time.Time `json:"updatedAt"`
	Models               []Model   `json:"models"`
}

// Catalog is the deployment-level hosted inventory. Organization eligibility
// is resolved separately.
type Catalog struct {
	Object    string    `json:"object"`
	UpdatedAt time.Time `json:"updatedAt"`
	Models    []Model   `json:"models"`
}

// Model is one hosted model's availability and pricing facts.
type Model struct {
	HostedModelID                             string                               `json:"hostedModelId"`
	ModelReference                            string                               `json:"modelReference"`
	DisplayName                               string                               `json:"displayName"`
	InventoryStatus                           string                               `json:"inventoryStatus"`
	EntitlementStatus                         string                               `json:"entitlementStatus"`
	AccessSource                              string                               `json:"accessSource,omitempty"`
	BalanceStatus                             string                               `json:"balanceStatus"`
	RuntimeStatus                             string                               `json:"runtimeStatus"`
	Enabled                                   bool                                 `json:"enabled"`
	Reasons                                   []string                             `json:"reasons,omitempty"`
	PricingFormulaVersion                     string                               `json:"pricingFormulaVersion,omitempty"`
	BillableDimensions                        []string                             `json:"billableDimensions,omitempty"`
	ContextTiers                              []PricingTier                        `json:"contextTiers,omitempty"`
	SupportedFamilies                         []string                             `json:"supportedFamilies,omitempty"`
	ExecutableWireAPIFamilies                 []string                             `json:"executableWireApiFamilies,omitempty"`
	MessagesFeatures                          []string                             `json:"messagesFeatures,omitempty"`
	ChatCompletionsFeatures                   []string                             `json:"chatCompletionsFeatures,omitempty"`
	ChatCompletionsCompatibilityProfile       *ChatCompletionsCompatibilityProfile `json:"chatCompletionsCompatibilityProfile,omitempty"`
	ResponsesFeatures                         []string                             `json:"responsesFeatures,omitempty"`
	InputCreditAtomsPerMillionTokens          *int64                               `json:"inputCreditAtomsPerMillionTokens,omitempty"`
	CacheReadInputCreditAtomsPerMillionTokens *int64                               `json:"cacheReadInputCreditAtomsPerMillionTokens,omitempty"`
	OutputCreditAtomsPerMillionTokens         *int64                               `json:"outputCreditAtomsPerMillionTokens,omitempty"`
}

// ChatCompletionsCompatibilityProfile is the versioned MaaS route contract
// used by AgentTeam server policy. The browser never supplies this profile.
type ChatCompletionsCompatibilityProfile struct {
	ID                                 string            `json:"id"`
	SupportedToolChoiceModes           []string          `json:"supported_tool_choice_modes"`
	ToolChoiceParameterAccepted        bool              `json:"tool_choice_parameter_accepted"`
	ParallelToolCallsParameterAccepted bool              `json:"parallel_tool_calls_parameter_accepted"`
	AssistantToolCallReasoningReplay   string            `json:"assistant_tool_call_reasoning_replay"`
	AssistantToolCallContent           string            `json:"assistant_tool_call_content"`
	OutputTokenField                   string            `json:"output_token_field"`
	SupportsStoreFalse                 bool              `json:"supports_store_false"`
	SupportsNOne                       bool              `json:"supports_n_one"`
	ReasoningEffortMap                 map[string]string `json:"reasoning_effort_map"`
}

var requiredChatCompletionsCompatibilityProfileFields = []string{
	"id",
	"supported_tool_choice_modes",
	"tool_choice_parameter_accepted",
	"parallel_tool_calls_parameter_accepted",
	"assistant_tool_call_reasoning_replay",
	"assistant_tool_call_content",
	"output_token_field",
	"supports_store_false",
	"supports_n_one",
	"reasoning_effort_map",
}

// UnmarshalJSON rejects partial or newer profile objects so an older AgentTeam
// binary cannot silently truncate MaaS compatibility semantics.
func (profile *ChatCompletionsCompatibilityProfile) UnmarshalJSON(data []byte) error {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(data, &fields); err != nil {
		return err
	}
	for _, field := range requiredChatCompletionsCompatibilityProfileFields {
		if _, ok := fields[field]; !ok {
			return fmt.Errorf("chat completions compatibility profile is missing field %q", field)
		}
	}

	type wireProfile ChatCompletionsCompatibilityProfile
	var decoded wireProfile
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&decoded); err != nil {
		return err
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		if err == nil {
			return fmt.Errorf("chat completions compatibility profile contains trailing JSON")
		}
		return err
	}
	*profile = ChatCompletionsCompatibilityProfile(decoded)
	return nil
}

// KnowledgeQAReasonLevels returns only caller-visible levels whose MaaS
// mapping preserves their meaning. Re-mapped low/medium levels are hidden.
func (profile ChatCompletionsCompatibilityProfile) KnowledgeQAReasonLevels() ([]string, string, bool) {
	if !profile.supportedKnowledgeQAContract() {
		return nil, "", false
	}
	levels := make([]string, 0, 3)
	for _, level := range []string{"low", "medium", "high"} {
		if strings.TrimSpace(profile.ReasoningEffortMap[level]) == level {
			levels = append(levels, level)
		}
	}
	if len(levels) == 0 {
		return nil, "", false
	}
	defaultLevel := levels[0]
	for _, level := range levels {
		if level == "medium" {
			defaultLevel = level
			break
		}
		if level == "high" {
			defaultLevel = level
		}
	}
	return levels, defaultLevel, true
}

// SupportsKnowledgeQAFeatures reports whether MaaS admits every independent
// Chat Completions parameter that Knowledge Q&A will send for this profile.
func (profile ChatCompletionsCompatibilityProfile) SupportsKnowledgeQAFeatures(features []string) bool {
	if !profile.supportedKnowledgeQAContract() ||
		!containsExactString(features, "chat_completions.tools") ||
		!containsExactString(features, "chat_completions.reasoning_effort") {
		return false
	}
	if profile.ToolChoiceParameterAccepted &&
		(!containsExactString(profile.SupportedToolChoiceModes, "required") ||
			!containsExactString(features, "chat_completions.tool_choice")) {
		return false
	}
	return !profile.ParallelToolCallsParameterAccepted ||
		containsExactString(features, "chat_completions.parallel_tool_calls")
}

func (profile ChatCompletionsCompatibilityProfile) supportedKnowledgeQAContract() bool {
	expected, ok := knownChatCompletionsCompatibilityProfile(profile.ID)
	if !ok {
		return false
	}
	actualJSON, actualErr := json.Marshal(profile)
	expectedJSON, expectedErr := json.Marshal(expected)
	return actualErr == nil && expectedErr == nil && string(actualJSON) == string(expectedJSON) &&
		containsExactString(profile.SupportedToolChoiceModes, "auto")
}

func knownChatCompletionsCompatibilityProfile(id string) (ChatCompletionsCompatibilityProfile, bool) {
	var profile ChatCompletionsCompatibilityProfile
	switch strings.TrimSpace(id) {
	case ChatCompletionsProfileOpenAIRequiredToolsV1:
		profile = ChatCompletionsCompatibilityProfile{
			ID:                                 ChatCompletionsProfileOpenAIRequiredToolsV1,
			SupportedToolChoiceModes:           []string{"none", "auto", "required", "named"},
			ToolChoiceParameterAccepted:        true,
			ParallelToolCallsParameterAccepted: true,
			AssistantToolCallReasoningReplay:   "optional",
			AssistantToolCallContent:           "nullable",
			OutputTokenField:                   "max_completion_tokens",
			SupportsStoreFalse:                 true,
			SupportsNOne:                       true,
			ReasoningEffortMap: map[string]string{
				"low": "low", "medium": "medium", "high": "high",
			},
		}
	case ChatCompletionsProfileDeepSeekV4ThinkingV1:
		profile = ChatCompletionsCompatibilityProfile{
			ID:                                 ChatCompletionsProfileDeepSeekV4ThinkingV1,
			SupportedToolChoiceModes:           []string{"auto"},
			ToolChoiceParameterAccepted:        false,
			ParallelToolCallsParameterAccepted: false,
			AssistantToolCallReasoningReplay:   "required",
			AssistantToolCallContent:           "non_null_string",
			OutputTokenField:                   "max_tokens",
			SupportsStoreFalse:                 false,
			SupportsNOne:                       false,
			ReasoningEffortMap: map[string]string{
				"low": "high", "medium": "high", "high": "high",
				"xhigh": "max", "max": "max",
			},
		}
	default:
		return ChatCompletionsCompatibilityProfile{}, false
	}
	return profile, true
}

func containsExactString(values []string, target string) bool {
	for _, value := range values {
		if strings.TrimSpace(value) == target {
			return true
		}
	}
	return false
}

// PricingTier is one MaaS context-length pricing tier for a hosted model.
type PricingTier struct {
	ID               string           `json:"id"`
	MinContextTokens int64            `json:"minContextTokens"`
	MaxContextTokens *int64           `json:"maxContextTokens,omitempty"`
	Rates            PricingTierRates `json:"rates"`
}

// PricingTierRates carries hosted credit rates for a context tier.
type PricingTierRates struct {
	InputTokensCreditAtomsPerMillionTokens    int64 `json:"inputTokensCreditAtomsPerMillionTokens"`
	CacheReadInputCreditAtomsPerMillionTokens int64 `json:"cacheReadInputCreditAtomsPerMillionTokens"`
	OutputCreditAtomsPerMillionTokens         int64 `json:"outputCreditAtomsPerMillionTokens"`
}

// EntitlementGrantRequest is the AgentTeam-to-agentteam-maas upsert payload.
type EntitlementGrantRequest struct {
	OrganizationID          string     `json:"organizationId"`
	HostedModelID           string     `json:"hostedModelId,omitempty"`
	HostedModelFamily       string     `json:"hostedModelFamily,omitempty"`
	CatalogPolicyHandle     string     `json:"catalogPolicyHandle,omitempty"`
	GrantSource             string     `json:"grantSource,omitempty"`
	DedicatedRoutePolicyRef string     `json:"dedicatedRoutePolicyRef,omitempty"`
	EffectiveFrom           *time.Time `json:"effectiveFrom,omitempty"`
	EffectiveTo             *time.Time `json:"effectiveTo,omitempty"`
	Status                  string     `json:"status,omitempty"`
}

// EntitlementGrant is the agentteam-maas response for one upserted hosted entitlement grant.
type EntitlementGrant struct {
	Object                  string     `json:"object"`
	ID                      string     `json:"id"`
	OrganizationID          string     `json:"organizationId"`
	HostedModelID           string     `json:"hostedModelId,omitempty"`
	HostedModelFamily       string     `json:"hostedModelFamily,omitempty"`
	CatalogPolicyHandle     string     `json:"catalogPolicyHandle,omitempty"`
	GrantSource             string     `json:"grantSource"`
	DedicatedRoutePolicyRef string     `json:"dedicatedRoutePolicyRef,omitempty"`
	EffectiveFrom           *time.Time `json:"effectiveFrom,omitempty"`
	EffectiveTo             *time.Time `json:"effectiveTo,omitempty"`
	Status                  string     `json:"status"`
	CreatedAt               time.Time  `json:"createdAt"`
	UpdatedAt               time.Time  `json:"updatedAt"`
}