package models

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

	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
)

type OpenAICompatibleLocalAdapter struct {
	base   *url.URL
	client *http.Client
	apiKey string
}

func NewOpenAICompatibleLocalAdapter(raw string, transport http.RoundTripper) (*OpenAICompatibleLocalAdapter, error) {
	return NewOpenAICompatibleLocalAdapterWithAPIKey(raw, "", transport)
}

// NewOpenAICompatibleLocalAdapterWithAPIKey creates an adapter for a gateway
// whose bearer credential was resolved by the trusted runtime boundary.
func NewOpenAICompatibleLocalAdapterWithAPIKey(raw, apiKey string, transport http.RoundTripper) (*OpenAICompatibleLocalAdapter, error) {
	return NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout(raw, apiKey, transport, 10*time.Second)
}

// NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout creates a local gateway
// adapter with an explicit invocation deadline. The request context remains
// authoritative for per-call cancellation; this client deadline prevents a
// provider from holding a worker indefinitely.
func NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout(raw, apiKey string, transport http.RoundTripper, timeout time.Duration) (*OpenAICompatibleLocalAdapter, error) {
	base, err := url.Parse(raw)
	if err != nil || base.Scheme != "http" || base.Host == "" || strings.TrimSuffix(base.Path, "/") != "/v1" || base.RawQuery != "" || base.Fragment != "" || base.User != nil {
		return nil, errors.New("invalid local compatible endpoint")
	}
	if timeout <= 0 {
		return nil, errors.New("local compatible timeout must be positive")
	}
	if transport == nil {
		transport = http.DefaultTransport
	}
	return &OpenAICompatibleLocalAdapter{base: base, client: &http.Client{Transport: transport, Timeout: timeout, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, apiKey: strings.TrimSpace(apiKey)}, nil
}

type compatibleModelsResponse struct {
	Data []struct {
		ID      string `json:"id"`
		Created int64  `json:"created"`
		OwnedBy string `json:"owned_by"`
		Meta    struct {
			Context         int64  `json:"n_ctx"`
			TrainingContext int64  `json:"n_ctx_train"`
			Parameters      int64  `json:"n_params"`
			Size            int64  `json:"size"`
			FileType        string `json:"ftype"`
		} `json:"meta"`
	} `json:"data"`
}

func (a *OpenAICompatibleLocalAdapter) Probe(ctx context.Context) ProviderProbe {
	_, probe := a.Discover(ctx)
	return probe
}

func (a *OpenAICompatibleLocalAdapter) Discover(ctx context.Context) ([]DiscoveredModel, ProviderProbe) {
	var response compatibleModelsResponse
	status, err := a.requestJSON(ctx, http.MethodGet, "/models", nil, &response)
	if err != nil {
		return nil, probeFailure(err)
	}
	if status != http.StatusOK {
		return nil, ProviderProbe{Health: capability.HealthUnavailable, Failure: mapCompatibleStatus(status)}
	}
	models := make([]DiscoveredModel, 0, len(response.Data))
	for _, item := range response.Data {
		if item.ID == "" {
			continue
		}
		stableMetadata := fmt.Sprintf("%s\x00%s\x00%d\x00%d\x00%d\x00%d\x00%s", item.ID, item.OwnedBy, item.Meta.Context, item.Meta.TrainingContext, item.Meta.Parameters, item.Meta.Size, item.Meta.FileType)
		digest := sha256.Sum256([]byte(stableMetadata))
		contextBytes := item.Meta.Context * 4
		models = append(models, DiscoveredModel{ProviderModelID: capability.ProviderModelID(item.ID), Fingerprint: capability.ModelFingerprint(hex.EncodeToString(digest[:])), MaxContextBytes: contextBytes, Modalities: []capability.ModelModality{capability.ModalityText}, Capabilities: []capability.ModelCapability{capability.CapabilityUnknown}, SizeBytes: item.Meta.Size, Runtime: item.OwnedBy, Quantization: item.Meta.FileType})
	}
	return models, ProviderProbe{Health: capability.HealthAvailable}
}

type compatibleChatRequest struct {
	Model       string              `json:"model"`
	Messages    []compatibleMessage `json:"messages"`
	Temperature int                 `json:"temperature"`
	Stream      bool                `json:"stream"`
	MaxTokens   int                 `json:"max_tokens"`
}
type compatibleMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}
type compatibleChatResponse struct {
	Choices []struct {
		FinishReason string            `json:"finish_reason"`
		Message      compatibleMessage `json:"message"`
	} `json:"choices"`
	Usage struct {
		Prompt     int64 `json:"prompt_tokens"`
		Completion int64 `json:"completion_tokens"`
	} `json:"usage"`
}

func (a *OpenAICompatibleLocalAdapter) Advise(ctx context.Context, request cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
	if err := ValidateAdvisoryRequest(request); err != nil || request.ProviderModelID == "" || len(request.Context) != 0 {
		return cognitive.AdvisoryResponse{}, ErrInvalidAdvisory
	}
	payload := compatibleChatRequest{Model: string(request.ProviderModelID), Messages: []compatibleMessage{{Role: "system", Content: "Provide non-canonical advisory text only. Never claim that an action occurred."}, {Role: "user", Content: request.Message}}, Temperature: 0, Stream: false, MaxTokens: request.MaxOutputBytes / 4}
	var result compatibleChatResponse
	status, err := a.requestJSON(ctx, http.MethodPost, "/chat/completions", payload, &result)
	if err != nil {
		return advisoryFailure(request, probeFailure(err).Failure), nil
	}
	if status != http.StatusOK {
		return advisoryFailure(request, mapCompatibleStatus(status)), nil
	}
	if len(result.Choices) != 1 || result.Choices[0].Message.Content == "" {
		return advisoryFailure(request, cognitive.FinishInvalidProviderResponse), nil
	}
	finish := cognitive.FinishCompleted
	if result.Choices[0].FinishReason == "length" {
		finish = cognitive.FinishLengthLimit
	}
	if finish != cognitive.FinishCompleted {
		return advisoryFailure(request, finish), nil
	}
	response := cognitive.AdvisoryResponse{Version: request.Version, RequestID: request.RequestID, ProviderBindingID: request.ProviderBindingID, ProviderBindingRevision: request.ProviderBindingRevision, AdvisoryText: result.Choices[0].Message.Content, FinishReason: finish, Usage: &cognitive.ProviderUsage{InputUnits: result.Usage.Prompt, OutputUnits: result.Usage.Completion}}
	if ValidateAdvisoryResponse(request, response) != nil {
		return advisoryFailure(request, cognitive.FinishInvalidProviderResponse), nil
	}
	return response, nil
}

func (a *OpenAICompatibleLocalAdapter) requestJSON(ctx context.Context, method, path string, input, output any) (int, error) {
	var body io.Reader
	if input != nil {
		encoded, err := json.Marshal(input)
		if err != nil {
			return 0, err
		}
		body = bytes.NewReader(encoded)
	}
	target := *a.base
	target.Path = strings.TrimSuffix(a.base.Path, "/") + path
	request, err := http.NewRequestWithContext(ctx, method, target.String(), body)
	if err != nil {
		return 0, err
	}
	request.Header.Set("Accept", "application/json")
	if a.apiKey != "" {
		request.Header.Set("Authorization", "Bearer "+a.apiKey)
	}
	if input != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	response, err := a.client.Do(request)
	if err != nil {
		return 0, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096))
		return response.StatusCode, nil
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, 2<<20))
	if err := decoder.Decode(output); err != nil {
		return response.StatusCode, fmt.Errorf("decode provider response: %w", err)
	}
	return response.StatusCode, nil
}

func probeFailure(err error) ProviderProbe {
	finish := cognitive.FinishProviderUnavailable
	if errors.Is(err, context.DeadlineExceeded) {
		finish = cognitive.FinishTimeout
	}
	var network net.Error
	if errors.As(err, &network) && network.Timeout() {
		finish = cognitive.FinishTimeout
	}
	return ProviderProbe{Health: capability.HealthUnavailable, Failure: finish}
}

func mapCompatibleStatus(status int) cognitive.FinishReason {
	switch status {
	case 401, 403:
		return cognitive.FinishAuthenticationFailed
	case 404:
		return cognitive.FinishModelUnavailable
	case 408, 504:
		return cognitive.FinishTimeout
	case 413:
		return cognitive.FinishContextTooLarge
	case 429:
		return cognitive.FinishRateLimited
	default:
		return cognitive.FinishProviderUnavailable
	}
}