/*
 * Copyright 2025 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package langsmith

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/bytedance/sonic"
)

// Langsmith func interface
type Langsmith interface {
	CreateRun(ctx context.Context, run *Run) error
	UpdateRun(ctx context.Context, runID string, patch *RunPatch) error
}

const (
	// DefaultLangsmithAPIURL Langsmith default url
	DefaultLangsmithAPIURL = "https://api.smith.langchain.com"
)

type RunType string

const (
	RunTypeChain RunType = "chain" // chain node
	RunTypeLLM   RunType = "llm"   // llm model node
	RunTypeTool  RunType = "tool"  // tool node
)

type Run struct {
	ID                 string                 `json:"id"`                             // Unique identifier for the span.
	Name               string                 `json:"name"`                           // The name associated with the run.
	RunType            RunType                `json:"run_type"`                       // Type of run, e.g., "llm", "chain", "tool".
	StartTime          time.Time              `json:"start_time"`                     // Start time of the run.
	EndTime            *time.Time             `json:"end_time,omitempty"`             // End time of the run.
	Inputs             map[string]interface{} `json:"inputs"`                         // A map or set of inputs provided to the run.
	Outputs            map[string]interface{} `json:"outputs,omitempty"`              // A map or set of outputs generated by the run.
	Error              *string                `json:"error,omitempty"`                // Error message if the run encountered an error.
	ParentRunID        *string                `json:"parent_run_id,omitempty"`        // Unique identifier of the parent run.
	TraceID            string                 `json:"trace_id,omitempty"`             // Unique identifier for the trace the run is a part of. This is also the id field of the root run of the trace
	Extra              map[string]interface{} `json:"extra,omitempty"`                // Any extra information run.
	SessionName        string                 `json:"session_name,omitempty"`         // langsmith session name
	ReferenceExampleID *string                `json:"reference_example_id,omitempty"` // ID of a reference example associated with the run. This is usually only present for evaluation runs.
	DottedOrder        string                 `json:"dotted_order,omitempty"`         // Ordering string, hierarchical. Format: run_start_timeZrun_uuid.child_run_start_timeZchild_run_uuid...
	Tags               []string               `json:"tags,omitempty"`                 // Tags or labels associated with the run.
}

// RunPatch update run when it is finished or failed, patch output or error msg.
type RunPatch struct {
	EndTime *time.Time             `json:"end_time,omitempty"` // End time of the run.
	Inputs  map[string]interface{} `json:"inputs,omitempty"`   // A map or set of inputs provided to the run.
	Outputs map[string]interface{} `json:"outputs,omitempty"`  // A map or set of outputs generated by the run.
	Error   *string                `json:"error,omitempty"`    // Error message if the run encountered an error.
	Extra   map[string]interface{} `json:"extra,omitempty"`    // Any extra information run.
}

type langsmithClient struct {
	apiKey     string
	baseURL    string
	httpClient *http.Client
}

// NewLangsmith create langsmith client
func NewLangsmith(apiKey, apiUrl string) Langsmith {
	if apiUrl == "" {
		apiUrl = DefaultLangsmithAPIURL
	}
	return &langsmithClient{
		apiKey:     apiKey,
		baseURL:    apiUrl,
		httpClient: &http.Client{Timeout: 10 * time.Second},
	}
}

// CreateRun create run
func (c *langsmithClient) CreateRun(ctx context.Context, run *Run) error {
	jsonData, err := sonic.Marshal(run)
	if err != nil {
		return fmt.Errorf("failed to marshal run data: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/runs", bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", c.apiKey)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to execute request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read response body: %w", err)
	}
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
		return fmt.Errorf("failed to create run, status: %s, body: %s", resp.Status, string(body))
	}

	// decode resp data
	err = sonic.Unmarshal(body, run)
	if err != nil {
		return fmt.Errorf("failed to decode response body: %w", err)
	}

	return nil
}

// UpdateRun update run when it is finished or failed, patch output or error msg.
func (c *langsmithClient) UpdateRun(ctx context.Context, runID string, patch *RunPatch) error {
	jsonData, err := json.Marshal(patch)
	if err != nil {
		return fmt.Errorf("failed to marshal patch data: %w", err)
	}

	url := fmt.Sprintf("%s/runs/%s", c.baseURL, runID)
	req, err := http.NewRequestWithContext(ctx, "PATCH", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", c.apiKey)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to execute request: %w", err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read response body: %w", err)
	}
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
		return fmt.Errorf("failed to update run, status: %s, body: %s", resp.Status, string(body))
	}

	return nil
}
