package memory

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/nativememory"
	"github.com/tidbcloud/AgentTeam/desktop/internal/provisioning/asset"
)

const maxNativeMemoryArchiveBytes = int64(68 << 20)
const maxNativeMemoryResponseBytes = int64(1 << 20)

type nativeRemote struct {
	conn     asset.ConnectionInfo
	fallback *http.Client
}

type nativeCheckpointFile struct {
	Path          string `json:"path"`
	ContentBase64 string `json:"contentBase64"`
	Executable    bool   `json:"executable,omitempty"`
}

type nativeCheckpointRequest struct {
	OperationID            string                 `json:"operationId"`
	ReplicaID              string                 `json:"replicaId"`
	ExpectedVersionID      string                 `json:"expectedVersionId"`
	Generation             uint64                 `json:"generation"`
	CoveredFromSequence    uint64                 `json:"coveredFromSequence"`
	CoveredThroughSequence uint64                 `json:"coveredThroughSequence"`
	Files                  []nativeCheckpointFile `json:"files"`
}

type nativeCheckpointResponse struct {
	VersionID              string `json:"versionId"`
	OperationID            string `json:"operationId"`
	ReplicaID              string `json:"replicaId"`
	Generation             uint64 `json:"generation"`
	CoveredFromSequence    uint64 `json:"coveredFromSequence"`
	CoveredThroughSequence uint64 `json:"coveredThroughSequence"`
}

func newNativeRemote(conn asset.ConnectionInfo, fallback *http.Client) *nativeRemote {
	return &nativeRemote{conn: conn, fallback: fallback}
}

func (r *nativeRemote) PublishCheckpoint(ctx context.Context, memoryID string, input nativememory.Checkpoint) (nativememory.Result, error) {
	files := make([]nativeCheckpointFile, 0, len(input.Files))
	for _, file := range input.Files {
		files = append(files, nativeCheckpointFile{
			Path: file.Path, ContentBase64: base64.StdEncoding.EncodeToString(file.Content), Executable: file.Executable,
		})
	}
	body, err := json.Marshal(nativeCheckpointRequest{
		OperationID: input.OperationID, ReplicaID: input.ReplicaID, ExpectedVersionID: input.ExpectedVersionID,
		Generation: input.Generation, CoveredFromSequence: input.CoveredFromSequence,
		CoveredThroughSequence: input.CoveredThroughSequence, Files: files,
	})
	if err != nil {
		return nativememory.Result{}, fmt.Errorf("memoryprovider: encode native checkpoint: %w", err)
	}
	var response nativeCheckpointResponse
	if err := r.doJSON(ctx, memoryID, body, &response); err != nil {
		return nativememory.Result{}, err
	}
	return nativememory.Result{
		VersionID: response.VersionID, OperationID: response.OperationID, ReplicaID: response.ReplicaID,
		Generation: response.Generation, CoveredFromSequence: response.CoveredFromSequence,
		CoveredThroughSequence: response.CoveredThroughSequence,
	}, nil
}

func (r *nativeRemote) OpenVersionArchive(ctx context.Context, memoryID, versionID string) (io.ReadCloser, error) {
	if r == nil {
		return nil, fmt.Errorf("memoryprovider: native remote is not configured")
	}
	serverURL := strings.TrimRight(strings.TrimSpace(r.conn.ServerURL), "/")
	workspaceID := strings.TrimSpace(r.conn.WorkspaceID)
	token := strings.TrimSpace(r.conn.UserToken)
	memoryID = strings.TrimSpace(memoryID)
	versionID = strings.TrimSpace(versionID)
	if serverURL == "" || workspaceID == "" || token == "" || memoryID == "" || versionID == "" {
		return nil, fmt.Errorf("memoryprovider: incomplete native archive connection")
	}
	endpoint := serverURL + "/api/v1/workspaces/" + url.PathEscape(workspaceID) + "/memory/" + url.PathEscape(memoryID) + "/versions/" + url.PathEscape(versionID) + "/archive"
	client := r.conn.HTTPClient
	if client == nil {
		client = r.fallback
	}
	if client == nil {
		client = http.DefaultClient
	}
	request := func(candidate string) (*http.Response, error) {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
		if err != nil {
			return nil, err
		}
		req.Header.Set("Authorization", "Bearer "+candidate)
		req.Header.Set("X-Workspace-Id", workspaceID)
		req.Header.Set("X-AgentTeam-Route-Workspace-Id", workspaceID)
		return client.Do(req)
	}
	response, err := request(token)
	if err != nil {
		return nil, fmt.Errorf("memoryprovider: open native archive: %w", err)
	}
	if response.StatusCode == http.StatusUnauthorized && r.conn.RefreshUserToken != nil {
		_ = response.Body.Close()
		refreshed, refreshErr := r.conn.RefreshUserToken(ctx, token)
		if refreshErr != nil {
			return nil, fmt.Errorf("memoryprovider: refresh native archive token: %w", refreshErr)
		}
		token = strings.TrimSpace(refreshed)
		if token == "" {
			return nil, fmt.Errorf("memoryprovider: refresh native archive token returned empty token")
		}
		response, err = request(token)
		if err != nil {
			return nil, fmt.Errorf("memoryprovider: open native archive after token refresh: %w", err)
		}
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("memoryprovider: open native archive: status %d", response.StatusCode)
	}
	data, err := io.ReadAll(io.LimitReader(response.Body, maxNativeMemoryArchiveBytes+1))
	if err != nil {
		return nil, fmt.Errorf("memoryprovider: read native archive: %w", err)
	}
	if int64(len(data)) > maxNativeMemoryArchiveBytes {
		return nil, fmt.Errorf("memoryprovider: native archive exceeds byte limit")
	}
	return io.NopCloser(bytes.NewReader(data)), nil
}

func (r *nativeRemote) doJSON(ctx context.Context, memoryID string, body []byte, out any) error {
	if r == nil {
		return fmt.Errorf("memoryprovider: native remote is not configured")
	}
	serverURL := strings.TrimRight(strings.TrimSpace(r.conn.ServerURL), "/")
	workspaceID := strings.TrimSpace(r.conn.WorkspaceID)
	token := strings.TrimSpace(r.conn.UserToken)
	if serverURL == "" || workspaceID == "" || token == "" || strings.TrimSpace(memoryID) == "" {
		return fmt.Errorf("memoryprovider: incomplete native remote connection")
	}
	endpoint := serverURL + "/api/v1/workspaces/" + url.PathEscape(workspaceID) + "/memory/" + url.PathEscape(strings.TrimSpace(memoryID)) + "/checkpoints"
	client := r.conn.HTTPClient
	if client == nil {
		client = r.fallback
	}
	if client == nil {
		client = http.DefaultClient
	}
	request := func(candidate string) (*http.Response, error) {
		req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
		if err != nil {
			return nil, err
		}
		req.Header.Set("Authorization", "Bearer "+candidate)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-Workspace-Id", workspaceID)
		req.Header.Set("X-AgentTeam-Route-Workspace-Id", workspaceID)
		return client.Do(req)
	}
	response, err := request(token)
	if err != nil {
		return fmt.Errorf("memoryprovider: publish native checkpoint: %w", err)
	}
	if response.StatusCode == http.StatusUnauthorized && r.conn.RefreshUserToken != nil {
		_ = response.Body.Close()
		refreshed, refreshErr := r.conn.RefreshUserToken(ctx, token)
		if refreshErr != nil {
			return fmt.Errorf("memoryprovider: refresh native checkpoint token: %w", refreshErr)
		}
		token = strings.TrimSpace(refreshed)
		if token == "" {
			return fmt.Errorf("memoryprovider: refresh native checkpoint token returned empty token")
		}
		response, err = request(token)
		if err != nil {
			return fmt.Errorf("memoryprovider: publish native checkpoint after token refresh: %w", err)
		}
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusCreated {
		return fmt.Errorf("memoryprovider: publish native checkpoint: status %d", response.StatusCode)
	}
	responseBody, err := io.ReadAll(io.LimitReader(response.Body, maxNativeMemoryResponseBytes+1))
	if err != nil {
		return fmt.Errorf("memoryprovider: read native checkpoint response: %w", err)
	}
	if int64(len(responseBody)) > maxNativeMemoryResponseBytes {
		return fmt.Errorf("memoryprovider: native checkpoint response exceeds byte limit")
	}
	if err := json.Unmarshal(responseBody, out); err != nil {
		return fmt.Errorf("memoryprovider: decode native checkpoint: %w", err)
	}
	return nil
}