// Copyright 2025 Microsoft Corporation
//
// 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 agent

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/anthropics/anthropic-sdk-go"
	"github.com/anthropics/anthropic-sdk-go/option"
	"github.com/anthropics/anthropic-sdk-go/vertex"
	"github.com/go-logr/logr"
)

const (
	// DefaultClaudeModel is the default Anthropic model used when no
	// model override is specified.
	DefaultClaudeModel = "claude-sonnet-4-20250514"

	// claudeMaxTokens is the maximum number of output tokens per API call.
	claudeMaxTokens int64 = 16384

	// ClaudeBackendAPI selects the direct Anthropic API backend (requires an API key).
	ClaudeBackendAPI = "api"

	// ClaudeBackendVertex selects Google Vertex AI as the backend
	// (uses Application Default Credentials for keyless auth).
	ClaudeBackendVertex = "vertex"
)

// Compile-time check: *ClaudeProvider implements LLMProvider.
var _ LLMProvider = (*ClaudeProvider)(nil)

// Compile-time check: *ClaudeSession implements LLMSession.
var _ LLMSession = (*ClaudeSession)(nil)

// ClaudeConfig holds configuration for the Anthropic Claude provider.
type ClaudeConfig struct {
	// APIKeyFile is the path to a file containing the Anthropic API key.
	// If empty, the ANTHROPIC_API_KEY environment variable is used.
	// Only used when Backend is "api" (the default).
	APIKeyFile string

	// Model is the Anthropic model to use (e.g. "claude-sonnet-4-20250514").
	// Defaults to DefaultClaudeModel.
	Model string

	// Backend selects the API backend: "api" for the direct Anthropic
	// API (default) or "vertex" for Google Vertex AI.
	Backend string

	// VertexProject is the GCP project ID. Required when Backend is "vertex".
	VertexProject string

	// VertexRegion is the GCP region (e.g. "us-east5"). Required when Backend is "vertex".
	VertexRegion string
}

// ClaudeProvider implements LLMProvider using the Anthropic Messages API.
// It manages a single anthropic.Client and creates ClaudeSession instances
// for individual analysis runs.
type ClaudeProvider struct {
	client anthropic.Client
	cfg    *ClaudeConfig
}

// NewClaudeProvider creates a ClaudeProvider configured with the given API key
// and model settings. When the Vertex AI backend is selected, ctx is used to
// initialise Google Application Default Credentials.
func NewClaudeProvider(ctx context.Context, cfg *ClaudeConfig) (*ClaudeProvider, error) {
	var opts []option.RequestOption

	switch cfg.Backend {
	case ClaudeBackendVertex:
		opts = append(opts, vertex.WithGoogleAuth(ctx, cfg.VertexRegion, cfg.VertexProject))
	default: // ClaudeBackendAPI or ""
		if cfg.APIKeyFile != "" {
			keyData, err := os.ReadFile(cfg.APIKeyFile)
			if err != nil {
				return nil, fmt.Errorf("reading Anthropic API key from %s: %w", cfg.APIKeyFile, err)
			}
			key := strings.TrimSpace(string(keyData))
			if key == "" {
				return nil, fmt.Errorf("anthropic API key file %s is empty", cfg.APIKeyFile)
			}
			opts = append(opts, option.WithAPIKey(key))
		}
		// If no key file, the SDK reads ANTHROPIC_API_KEY from the environment.
	}

	return &ClaudeProvider{
		client: anthropic.NewClient(opts...),
		cfg:    cfg,
	}, nil
}

// CreateProviderSession creates a new ClaudeSession for an analysis run.
// It concatenates cfg.IdentityPrompt, cfg.TonePrompt, and cfg.SystemPrompt
// into a single system message. All three fields are set centrally by the
// caller (e.g. analyze_cmd.go).
func (p *ClaudeProvider) CreateProviderSession(ctx context.Context, logger logr.Logger, cfg ProviderSessionConfig) (LLMSession, error) {
	model := cfg.Model
	if model == "" {
		model = p.cfg.Model
	}
	if model == "" {
		model = DefaultClaudeModel
	}

	// Convert provider-neutral tool definitions to Anthropic tool params.
	tools := make([]anthropic.ToolUnionParam, 0, len(cfg.Tools))
	for _, td := range cfg.Tools {
		at, err := toolDefinitionToAnthropicTool(td)
		if err != nil {
			return nil, fmt.Errorf("converting tool %q: %w", td.Name, err)
		}
		tools = append(tools, at)
	}

	// Build the full system prompt from the three parts passed in cfg.
	// The caller is responsible for setting IdentityPrompt, TonePrompt,
	// and SystemPrompt — this provider just concatenates them.
	fullSystemPrompt := cfg.IdentityPrompt + "\n\n" + cfg.TonePrompt + "\n\n" + cfg.SystemPrompt

	sessionID := fmt.Sprintf("claude-%d", time.Now().UnixNano())
	logger = logger.WithValues("sessionID", sessionID)
	logger.Info("Created Claude session.", "model", model)

	return &ClaudeSession{
		client:       p.client,
		model:        model,
		systemPrompt: fullSystemPrompt,
		tools:        tools,
		toolHandlers: buildToolHandlerMap(cfg.Tools),
		messages:     nil,
		sessionID:    sessionID,
		logger:       logger,
	}, nil
}

// Stop is a no-op for the Claude provider — the HTTP client has no
// long-running subprocess to shut down.
func (p *ClaudeProvider) Stop() error {
	return nil
}

// ClaudeSession implements LLMSession using the Anthropic Messages API.
// It maintains the conversation history and handles the tool-use loop
// internally within SendAndWait.
type ClaudeSession struct {
	client       anthropic.Client
	model        string
	systemPrompt string
	tools        []anthropic.ToolUnionParam
	toolHandlers map[string]func(ctx context.Context, params json.RawMessage) (string, error)
	// messages is the active API context sent to the model on each turn.
	// ResetHistory clears this to bound context growth between phases.
	messages []anthropic.MessageParam
	// archivedMessages is the full, append-only transcript of the session.
	// It is never cleared by ResetHistory so that SaveConversation can
	// persist the complete conversation for later inspection.
	archivedMessages []anthropic.MessageParam
	sessionID        string
	logger           logr.Logger
}

// SessionID returns the unique identifier for this session.
func (s *ClaudeSession) SessionID() string {
	return s.sessionID
}

// SendAndWait sends a user prompt and blocks until Claude finishes responding,
// including any tool-use rounds. It implements the tool-use loop: when Claude
// returns tool_use content blocks, the session calls the corresponding handler,
// sends the tool_result back, and waits for the next response.
func (s *ClaudeSession) SendAndWait(ctx context.Context, prompt string) (string, error) {
	s.logger.V(1).Info("Sending message to Claude session.", "promptLength", len(prompt))

	// Add user message to both the active API context and the archived
	// transcript. archivedMessages survives ResetHistory so the full
	// conversation can still be persisted by SaveConversation.
	userMsg := anthropic.NewUserMessage(anthropic.NewTextBlock(prompt))
	s.messages = append(s.messages, userMsg)
	s.archivedMessages = append(s.archivedMessages, userMsg)

	// Tool-use loop.
	for {
		params := anthropic.MessageNewParams{
			Model:     s.model,
			MaxTokens: claudeMaxTokens,
			// Mark the system prompt as an ephemeral cache breakpoint. The
			// system prompt is identical across every turn, so Anthropic
			// prompt caching lets subsequent requests reuse it, cutting
			// input-token cost and latency for repeated calls.
			System: []anthropic.TextBlockParam{
				{
					Text: s.systemPrompt,
					CacheControl: anthropic.CacheControlEphemeralParam{
						TTL: anthropic.CacheControlEphemeralTTLTTL5m,
					},
				},
			},
			Messages: s.messages,
		}
		if len(s.tools) > 0 {
			params.Tools = s.tools
		}

		s.logger.V(3).Info("Calling Anthropic Messages API.", "messageCount", len(s.messages))

		resp, err := s.client.Messages.New(ctx, params)
		if err != nil {
			return "", fmt.Errorf("claude API call failed: %w", err)
		}

		s.logger.V(3).Info("Claude API responded.",
			"stopReason", resp.StopReason,
			"inputTokens", resp.Usage.InputTokens,
			"outputTokens", resp.Usage.OutputTokens,
		)

		// Add assistant response to both the active context and the archive.
		assistantMsg := resp.ToParam()
		s.messages = append(s.messages, assistantMsg)
		s.archivedMessages = append(s.archivedMessages, assistantMsg)

		// If stop reason is not tool_use, extract text and return.
		if resp.StopReason != anthropic.StopReasonToolUse {
			return extractTextFromResponse(resp), nil
		}

		// Process tool use blocks.
		toolResults, err := s.processToolUseBlocks(ctx, resp)
		if err != nil {
			return "", err
		}

		// Add tool results as a user message to both the active context and
		// the archive, then continue the loop.
		toolResultMsg := anthropic.NewUserMessage(toolResults...)
		s.messages = append(s.messages, toolResultMsg)
		s.archivedMessages = append(s.archivedMessages, toolResultMsg)
	}
}

// processToolUseBlocks iterates over the response content blocks, identifies
// tool_use blocks, calls the corresponding handler, and returns tool_result
// content blocks.
func (s *ClaudeSession) processToolUseBlocks(ctx context.Context, resp *anthropic.Message) ([]anthropic.ContentBlockParamUnion, error) {
	var results []anthropic.ContentBlockParamUnion

	for _, block := range resp.Content {
		variant := block.AsAny()
		toolUse, ok := variant.(anthropic.ToolUseBlock)
		if !ok {
			continue
		}

		s.logger.V(3).Info("Claude requested tool call.",
			"toolName", toolUse.Name,
			"toolCallID", toolUse.ID,
		)

		handler, exists := s.toolHandlers[toolUse.Name]
		if !exists {
			errMsg := fmt.Sprintf("unknown tool %q", toolUse.Name)
			s.logger.Error(fmt.Errorf("unknown tool: %s", toolUse.Name), "Tool not found.")
			results = append(results, anthropic.NewToolResultBlock(
				toolUse.ID, errMsg, true,
			))
			continue
		}

		// Extract raw JSON input from the tool use block.
		rawInput := toolUse.Input

		result, err := handler(ctx, rawInput)
		if err != nil {
			s.logger.V(3).Info("Tool call returned error.",
				"toolName", toolUse.Name,
				"error", err.Error(),
			)
			results = append(results, anthropic.NewToolResultBlock(
				toolUse.ID, err.Error(), true,
			))
			continue
		}

		s.logger.V(3).Info("Tool call succeeded.",
			"toolName", toolUse.Name,
			"resultLength", len(result),
		)
		results = append(results, anthropic.NewToolResultBlock(
			toolUse.ID, result, false,
		))
	}

	return results, nil
}

// ResetHistory clears the active conversation context. This is used
// between analysis phases (e.g. before review rounds) to prevent the
// accumulated context from exceeding the model's token limit. Only the
// active context (s.messages) is cleared; the archived transcript
// (s.archivedMessages) is preserved so SaveConversation can still persist
// the complete conversation. The system prompt, tools, and tool handlers
// are also preserved.
func (s *ClaudeSession) ResetHistory() {
	s.logger.V(1).Info("Resetting Claude active conversation context; archived transcript preserved.",
		"clearedMessageCount", len(s.messages),
		"archivedMessageCount", len(s.archivedMessages),
	)
	s.messages = nil
}

// SaveConversation writes the full archived conversation transcript to a
// JSON file. It persists s.archivedMessages (the append-only transcript)
// rather than s.messages (the active context), so the complete
// conversation is saved even after ResetHistory has trimmed the active
// context between phases.
func (s *ClaudeSession) SaveConversation(path string) {
	if len(s.archivedMessages) == 0 {
		s.logger.Info("No conversation messages to save.")
		return
	}

	data, err := json.MarshalIndent(s.archivedMessages, "", "  ")
	if err != nil {
		s.logger.Error(err, "Failed to marshal conversation for saving.")
		return
	}

	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		s.logger.Error(err, "Failed to create directory for conversation dump.", "path", path)
		return
	}

	if err := os.WriteFile(path, data, 0644); err != nil {
		s.logger.Error(err, "Failed to write conversation to disk.", "path", path)
		return
	}

	s.logger.Info("Wrote conversation to disk.", "path", path, "size", len(data))
}

// Disconnect is a no-op for the Claude provider — there is no persistent
// session state to disconnect from.
func (s *ClaudeSession) Disconnect() error {
	return nil
}

// Delete is a no-op for the Claude provider — there is no server-side
// session state to delete.
func (s *ClaudeSession) Delete(_ context.Context) error {
	return nil
}

// toolDefinitionToAnthropicTool converts a provider-neutral ToolDefinition
// to an Anthropic ToolUnionParam.
func toolDefinitionToAnthropicTool(td ToolDefinition) (anthropic.ToolUnionParam, error) {
	// Parse the JSON schema to extract properties and required fields.
	var schema struct {
		Properties map[string]any `json:"properties"`
		Required   []string       `json:"required"`
	}
	if len(td.ParamSchema) > 0 {
		if err := json.Unmarshal(td.ParamSchema, &schema); err != nil {
			return anthropic.ToolUnionParam{}, fmt.Errorf("parsing parameter schema: %w", err)
		}
	}

	return anthropic.ToolUnionParam{
		OfTool: &anthropic.ToolParam{
			Name:        td.Name,
			Description: anthropic.String(td.Description),
			InputSchema: anthropic.ToolInputSchemaParam{
				Properties: schema.Properties,
				Required:   schema.Required,
			},
		},
	}, nil
}

// buildToolHandlerMap creates a name→handler lookup map from tool definitions.
func buildToolHandlerMap(tools []ToolDefinition) map[string]func(ctx context.Context, params json.RawMessage) (string, error) {
	m := make(map[string]func(ctx context.Context, params json.RawMessage) (string, error), len(tools))
	for _, td := range tools {
		m[td.Name] = td.Handler
	}
	return m
}

// extractTextFromResponse concatenates all text content blocks from an
// Anthropic response into a single string.
func extractTextFromResponse(resp *anthropic.Message) string {
	var parts []string
	for _, block := range resp.Content {
		variant := block.AsAny()
		if textBlock, ok := variant.(anthropic.TextBlock); ok {
			parts = append(parts, textBlock.Text)
		}
	}
	return strings.Join(parts, "")
}
