/*
 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
 * openFuyao is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package rdma

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/openfuyao/weight-dispatcher/pkg/internal/errutil"
	sharedtypes "github.com/openfuyao/weight-dispatcher/pkg/types"
)

// HuggingFaceChunkClient reads model files directly from HuggingFace endpoints.
type HuggingFaceChunkClient struct {
	client *http.Client
	token  string
}

// NewHuggingFaceChunkClientWithToken creates a chunk client for HuggingFace-hosted artifacts.
func NewHuggingFaceChunkClientWithToken(client *http.Client, token string) *HuggingFaceChunkClient {
	if client == nil {
		client = &http.Client{
			Timeout:   120 * time.Second,
			Transport: newHuggingFaceTransport(),
		}
	}
	return &HuggingFaceChunkClient{client: client, token: token}
}

// Stat returns one HuggingFace file size via a HEAD request.
func (c *HuggingFaceChunkClient) Stat(ctx context.Context, endpoint, rootPath, relativePath string) (_ int64, err error) {

	resolution := resolveHFURL(endpoint, rootPath, relativePath)
	req, err := http.NewRequestWithContext(ctx, http.MethodHead, resolution.url, http.NoBody)
	if err != nil {
		return 0, fmt.Errorf("hf stat build request: %w", err)
	}
	if tok := firstNonEmpty(c.token, resolution.token); tok != "" {
		req.Header.Set("Authorization", "Bearer "+tok)
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return 0, fmt.Errorf("hf stat %s: %w", resolution.url, err)
	}
	defer func() {
		err = mergeCloseError(err, resp.Body.Close(), fmt.Sprintf("close hf stat response body %s", resolution.url))
	}()
	if resp.StatusCode == http.StatusNotFound {
		return 0, fmt.Errorf("hf file not found: %s", resolution.url)
	}
	if resp.StatusCode >= 300 {
		return 0, fmt.Errorf("hf stat returned status %d for %s", resp.StatusCode, resolution.url)
	}
	if resp.ContentLength < 0 {
		return 0, fmt.Errorf("hf stat: content-length not returned for %s", resolution.url)
	}
	return resp.ContentLength, nil
}

// ReadAt reads one byte range from a HuggingFace-hosted file.
func (c *HuggingFaceChunkClient) ReadAt(ctx context.Context, endpoint, rootPath, relativePath string, offset, length int64) (_ []byte, err error) {
	resolution := resolveHFURL(endpoint, rootPath, relativePath)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, resolution.url, http.NoBody)
	if err != nil {
		return nil, fmt.Errorf("hf readat build request: %w", err)
	}
	if tok := firstNonEmpty(c.token, resolution.token); tok != "" {
		req.Header.Set("Authorization", "Bearer "+tok)
	}
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", offset, offset+length-1))
	resp, err := c.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("hf readat %s [%d,%d): %w", resolution.url, offset, offset+length, err)
	}
	defer func() {
		err = mergeCloseError(err, resp.Body.Close(), fmt.Sprintf("close hf readat response body %s", resolution.url))
	}()
	if resp.StatusCode != http.StatusPartialContent && resp.StatusCode != http.StatusOK {
		body, readErr := io.ReadAll(io.LimitReader(resp.Body, 512))
		if readErr != nil {
			return nil, fmt.Errorf("hf readat %s read error response body: %w", resolution.url, readErr)
		}
		return nil, fmt.Errorf("hf readat %s returned status %d: %s", resolution.url, resp.StatusCode, string(body))
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, length+1))
	if err != nil {
		return nil, fmt.Errorf("hf readat %s: %w", resolution.url, err)
	}
	if int64(len(data)) != length {
		return nil, fmt.Errorf("hf readat %s: expected %d bytes, got %d", resolution.url, length, len(data))
	}
	return data, nil
}

func resolveHFURL(endpoint, modelID, relativePath string) hfURLResolution {
	base := strings.TrimRight(endpoint, "/")
	if base == "" {
		base = "https://huggingface.co"
	}
	revision := ExtractHFRevision(base)
	token := ""
	if idx := strings.Index(base, "#rev="); idx >= 0 {
		base = base[:idx]
	}
	if idx := strings.Index(base, "?token="); idx >= 0 {
		token = base[idx+7:]
		base = base[:idx]
	}
	return hfURLResolution{
		url:      fmt.Sprintf("%s/%s/resolve/%s/%s", base, strings.TrimLeft(modelID, "/"), revision, strings.TrimLeft(relativePath, "/")),
		token:    token,
		revision: revision,
	}
}

// ExtractHFRevision returns the explicit revision encoded in a HuggingFace endpoint.
func ExtractHFRevision(endpoint string) string {
	if idx := strings.Index(endpoint, "#rev="); idx >= 0 {
		return endpoint[idx+5:]
	}
	return "main"
}

func firstNonEmpty(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

type hfModelFile struct {
	RFileName string     `json:"rfilename"`
	Size      int64      `json:"size"`
	BlobID    string     `json:"blobId,omitempty"`
	LFS       *hfLFSInfo `json:"lfs,omitempty"`
}

type hfLFSInfo struct {
	OID string `json:"oid,omitempty"`
}

type hfModelInfo struct {
	Sha      string        `json:"sha,omitempty"`
	Siblings []hfModelFile `json:"siblings"`
}

type hfURLResolution struct {
	url      string
	token    string
	revision string
}

type huggingFaceManifestEndpoint struct {
	endpoint string
	token    string
}

// HuggingFaceManifestRequest describes one logical-manifest lookup against the HuggingFace model API.
type HuggingFaceManifestRequest struct {
	Endpoint       string
	ModelID        string
	Token          string
	Revision       string
	ChunkSizeBytes int64
}

var hfBlobSHA256Pattern = regexp.MustCompile(`^[a-f0-9]{64}$`)

const defaultHuggingFaceManifestChunkSizeBytes int64 = 64 * 1024 * 1024

// ResolveHuggingFaceManifest builds a logical manifest from the HuggingFace model API.
func ResolveHuggingFaceManifest(
	ctx context.Context,
	httpClient *http.Client,
	req HuggingFaceManifestRequest,
) (sharedtypes.LogicalManifest, error) {
	httpClient = ensureHuggingFaceManifestHTTPClient(httpClient)
	normalizedEndpoint := normalizeHuggingFaceManifestEndpoint(req.Endpoint, req.Token)
	revision := defaultHuggingFaceRevision(req.Revision)
	chunkSizeBytes := normalizeHuggingFaceManifestChunkSize(req.ChunkSizeBytes)
	apiURL := huggingFaceManifestAPIURL(normalizedEndpoint.endpoint, req.ModelID)
	info, err := fetchHuggingFaceModelInfo(ctx, httpClient, apiURL, normalizedEndpoint.token)
	if err != nil {
		return sharedtypes.LogicalManifest{}, err
	}
	return buildHuggingFaceManifest(req.ModelID, revision, chunkSizeBytes, info), nil
}

func ensureHuggingFaceManifestHTTPClient(httpClient *http.Client) *http.Client {
	if httpClient != nil {
		return httpClient
	}
	return &http.Client{
		Timeout:   30 * time.Second,
		Transport: newHuggingFaceTransport(),
	}
}

func normalizeHuggingFaceManifestEndpoint(endpoint, token string) huggingFaceManifestEndpoint {
	if endpoint == "" {
		endpoint = "https://huggingface.co"
	}
	if idx := strings.Index(endpoint, "#rev="); idx >= 0 {
		endpoint = endpoint[:idx]
	}
	if idx := strings.Index(endpoint, "?token="); idx >= 0 {
		if token == "" {
			token = endpoint[idx+7:]
		}
		endpoint = endpoint[:idx]
	}
	return huggingFaceManifestEndpoint{
		endpoint: endpoint,
		token:    token,
	}
}

func defaultHuggingFaceRevision(revision string) string {
	if revision != "" {
		return revision
	}
	return "main"
}

func normalizeHuggingFaceManifestChunkSize(chunkSizeBytes int64) int64 {
	if chunkSizeBytes > 0 {
		return chunkSizeBytes
	}
	return defaultHuggingFaceManifestChunkSizeBytes
}

func huggingFaceManifestAPIURL(endpoint, modelID string) string {
	return fmt.Sprintf("%s/api/models/%s?blobs=true", strings.TrimRight(endpoint, "/"), modelID)
}

func fetchHuggingFaceModelInfo(
	ctx context.Context,
	httpClient *http.Client,
	apiURL, token string,
) (_ hfModelInfo, err error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, http.NoBody)
	if err != nil {
		return hfModelInfo{}, fmt.Errorf("build hf api request: %w", err)
	}
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	resp, err := httpClient.Do(req)
	if err != nil {
		return hfModelInfo{}, fmt.Errorf("hf api request: %w", err)
	}
	defer func() {
		err = mergeCloseError(err, resp.Body.Close(), fmt.Sprintf("close hf api response body %s", apiURL))
	}()
	if err := validateHuggingFaceAPIResponse(resp); err != nil {
		return hfModelInfo{}, err
	}
	return decodeHuggingFaceModelInfo(resp.Body)
}

func validateHuggingFaceAPIResponse(resp *http.Response) error {
	if resp.StatusCode < 300 {
		return nil
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, 512))
	if err != nil {
		return errutil.Wrap("read hf api error response body", err)
	}
	return fmt.Errorf("hf api returned status %d: %s", resp.StatusCode, string(body))
}

func decodeHuggingFaceModelInfo(body io.Reader) (hfModelInfo, error) {
	var info hfModelInfo
	if err := json.NewDecoder(body).Decode(&info); err != nil {
		return hfModelInfo{}, fmt.Errorf("decode hf api response: %w", err)
	}
	return info, nil
}

func buildHuggingFaceManifest(
	modelID, revision string,
	chunkSizeBytes int64,
	info hfModelInfo,
) sharedtypes.LogicalManifest {
	return sharedtypes.LogicalManifest{
		ArtifactKey:    fmt.Sprintf("hf://%s@%s", modelID, revision),
		RootPath:       modelID,
		ChunkSizeBytes: chunkSizeBytes,
		Digest:         info.Sha,
		Files:          buildHuggingFaceManifestFiles(info),
	}
}

func buildHuggingFaceManifestFiles(info hfModelInfo) []sharedtypes.ArtifactFile {
	files := make([]sharedtypes.ArtifactFile, 0, len(info.Siblings))
	for _, sibling := range info.Siblings {
		if sibling.RFileName == "" {
			continue
		}
		files = append(files, buildHuggingFaceArtifactFile(info.Sha, sibling))
	}
	return files
}

func buildHuggingFaceArtifactFile(commitHash string, file hfModelFile) sharedtypes.ArtifactFile {
	kind := classifyHFFile(file.RFileName)
	sha256 := resolveHFFileSHA256(file)
	return sharedtypes.ArtifactFile{
		RelativePath: file.RFileName,
		SizeBytes:    file.Size,
		Kind:         kind,
		Chunkable:    kind == sharedtypes.ArtifactFileKindSafeTensors,
		Required:     true,
		ETag:         resolveHFFileETag(file, sha256),
		SHA256:       sha256,
		CommitHash:   commitHash,
	}
}

func resolveHFFileSHA256(file hfModelFile) string {
	if file.LFS != nil && file.LFS.OID != "" {
		return strings.ToLower(strings.TrimSpace(file.LFS.OID))
	}
	base := path.Base(strings.TrimSpace(file.RFileName))
	if hfBlobSHA256Pattern.MatchString(strings.ToLower(base)) {
		return strings.ToLower(base)
	}
	return ""
}

func resolveHFFileETag(file hfModelFile, sha256 string) string {
	if sha256 != "" {
		return sha256
	}
	if file.BlobID != "" {
		return strings.TrimSpace(file.BlobID)
	}
	return ""
}

func classifyHFFile(name string) sharedtypes.ArtifactFileKind {
	lower := strings.ToLower(name)
	switch {
	case strings.HasSuffix(lower, ".safetensors"):
		return sharedtypes.ArtifactFileKindSafeTensors
	case strings.HasSuffix(lower, ".json") || strings.HasSuffix(lower, ".yaml") || strings.HasSuffix(lower, ".yml"):
		return sharedtypes.ArtifactFileKindJSON
	case strings.Contains(lower, "tokenizer") || strings.HasSuffix(lower, ".model") || strings.HasSuffix(lower, ".vocab"):
		return sharedtypes.ArtifactFileKindTokenizer
	default:
		return sharedtypes.ArtifactFileKindAuxiliary
	}
}

// IsHuggingFaceEndpoint reports whether one source should use the HuggingFace client.
func IsHuggingFaceEndpoint(sourceType, endpoint string) bool {
	if strings.EqualFold(sourceType, "huggingface") || strings.EqualFold(sourceType, "hf") {
		return true
	}
	return strings.Contains(endpoint, "huggingface.co") || strings.Contains(endpoint, "hf.co")
}

func newHuggingFaceTransport() *http.Transport {
	defaultTransport, ok := http.DefaultTransport.(*http.Transport)
	if !ok {
		return &http.Transport{
			ForceAttemptHTTP2:   true,
			MaxIdleConns:        256,
			MaxIdleConnsPerHost: 64,
			MaxConnsPerHost:     64,
			IdleConnTimeout:     90 * time.Second,
		}
	}
	base := defaultTransport.Clone()
	base.ForceAttemptHTTP2 = true
	base.MaxIdleConns = 256
	base.MaxIdleConnsPerHost = 64
	base.MaxConnsPerHost = 64
	base.IdleConnTimeout = 90 * time.Second
	return base
}