package cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/pkg/browser"
	"github.com/spf13/cobra"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/tokenanalysis"
)

type tokenAnalyzeOptions struct {
	window            string
	agentRef          string
	runRef            string
	format            string
	output            string
	caseID            string
	changedVariable   string
	workspaceRevision string
	toolProfile       string
	cacheState        string
	outcome           string
	stop              bool
}

type tokenCompareOptions struct {
	baseline  string
	candidate string
	format    string
}

var tokenCmd = newTokenCommand()
var newTokenAnalysisClient = tokenAnalysisClient

func newTokenCommand() *cobra.Command {
	command := &cobra.Command{
		Use:   "token",
		Short: "Diagnose AgentTeam Token consumption",
		Args:  cobra.NoArgs,
	}
	command.AddCommand(newTokenAnalyzeCommand(), newTokenCompareCommand(), newTokenUICommand())
	return command
}

func newTokenAnalyzeCommand() *cobra.Command {
	options := &tokenAnalyzeOptions{}
	command := &cobra.Command{
		Use:   "analyze",
		Short: "Analyze local AgentTeam Token usage and Context",
		Args:  cobra.NoArgs,
		RunE: func(command *cobra.Command, _ []string) error {
			return runTokenAnalyze(command, options)
		},
	}
	flags := command.Flags()
	flags.StringVar(&options.window, "window", "1d", "analysis window (6h|12h|1d|24h|3d|1w)")
	flags.StringVar(&options.agentRef, "agent", "", "analyze one opaque Agent ref")
	flags.StringVar(&options.runRef, "run", "", "analyze one opaque Run ref")
	flags.StringVar(&options.caseID, "case", "", "controlled experiment case identifier")
	flags.StringVar(&options.changedVariable, "changed-variable", "", "single changed variable, for example platform.agentteam.exec")
	flags.StringVar(&options.workspaceRevision, "workspace-revision", "", "controlled Workspace/repository revision identifier")
	flags.StringVar(&options.toolProfile, "tool-profile", "", "controlled tool and permission profile identifier")
	flags.StringVar(&options.cacheState, "cache-state", "unknown", "cache state (cold|warm|unknown)")
	flags.StringVar(&options.outcome, "outcome", "unknown", "task outcome (success|failure|unknown)")
	flags.StringVarP(&options.format, "format", "f", "table", "output format (table|json)")
	flags.StringVarP(&options.output, "output", "o", "", "write a privacy-safe report JSON")
	flags.BoolVar(&options.stop, "stop", false, "report current evidence, then stop detailed collection")
	return command
}

func newTokenCompareCommand() *cobra.Command {
	options := &tokenCompareOptions{}
	command := &cobra.Command{
		Use:   "compare",
		Short: "Compare two saved AgentTeam Token reports",
		Args:  cobra.NoArgs,
		RunE: func(command *cobra.Command, _ []string) error {
			if strings.TrimSpace(options.baseline) == "" || strings.TrimSpace(options.candidate) == "" {
				return errors.New("--baseline and --candidate are required")
			}
			baseline, err := tokenanalysis.ReadReport(options.baseline)
			if err != nil {
				return err
			}
			candidate, err := tokenanalysis.ReadReport(options.candidate)
			if err != nil {
				return err
			}
			return renderTokenComparison(command, tokenanalysis.Compare(baseline, candidate), options.format)
		},
	}
	flags := command.Flags()
	flags.StringVar(&options.baseline, "baseline", "", "baseline report JSON")
	flags.StringVar(&options.candidate, "candidate", "", "candidate report JSON")
	flags.StringVarP(&options.format, "format", "f", "table", "output format (table|json)")
	return command
}

func newTokenUICommand() *cobra.Command {
	var noOpen bool
	command := &cobra.Command{
		Use:   "ui",
		Short: "Open the local AgentTeam Token analysis UI",
		Args:  cobra.NoArgs,
		RunE: func(command *cobra.Command, _ []string) error {
			client, agentCaller, err := newTokenAnalysisClient()
			if err != nil {
				return err
			}
			if agentCaller {
				return errors.New("Agent caller cannot open the machine-wide Token UI")
			}
			// UI sessions may outlive the five-minute CLI credential lifetime.
			// Mint a fresh stateless credential for each daemon request.
			client.APIKeyProvider = token.Default().GenerateTempKey
			ctx, stop := signal.NotifyContext(command.Context(), os.Interrupt)
			defer stop()
			return tokenanalysis.RunUI(ctx, tokenanalysis.UIOptions{
				Client:  client,
				NoOpen:  noOpen,
				OpenURL: browser.OpenURL,
				Out:     command.OutOrStdout(),
			})
		},
	}
	command.Flags().BoolVar(&noOpen, "no-open", false, "print the access URL without opening a browser")
	return command
}

func runTokenAnalyze(command *cobra.Command, options *tokenAnalyzeOptions) error {
	if options == nil {
		return errors.New("analyze options are required")
	}
	if _, err := tokenanalysis.ParseWindow(options.window); err != nil {
		return err
	}
	if strings.TrimSpace(options.changedVariable) != "" && strings.TrimSpace(options.runRef) == "" {
		return errors.New("--changed-variable requires --run")
	}
	client, agentCaller, err := newTokenAnalysisClient()
	if err != nil {
		return err
	}
	if agentCaller && options.stop {
		return errors.New("Agent caller cannot change Token collection")
	}
	ctx, cancel := context.WithTimeout(command.Context(), 20*time.Second)
	defer cancel()
	mode := tokenanalysis.AnalysisModeDefault
	if options.stop {
		mode = tokenanalysis.AnalysisModeStop
	}
	report, err := client.Analyze(ctx, tokenanalysis.Request{
		Window:   options.window,
		AgentRef: options.agentRef,
		RunRef:   options.runRef,
		Experiment: tokenanalysis.ExperimentMetadata{
			Case:            options.caseID,
			ChangedVariable: options.changedVariable,
			Workspace:       options.workspaceRevision,
			Tools:           options.toolProfile,
			CacheState:      options.cacheState,
			Outcome:         options.outcome,
		},
	}, mode)
	if err != nil {
		return err
	}
	if strings.TrimSpace(options.output) != "" {
		if err := writeTokenReport(options.output, tokenanalysis.Portable(report)); err != nil {
			return err
		}
	}
	return renderTokenReport(command, report, options.format)
}

func tokenAnalysisClient() (tokenanalysis.Client, bool, error) {
	agentToken := strings.TrimSpace(os.Getenv(tokenanalysis.EnvToken))
	if agentToken != "" {
		baseURL := strings.TrimSpace(os.Getenv(tokenanalysis.EnvBaseURL))
		agentID := strings.TrimSpace(os.Getenv(tokenanalysis.EnvAgentID))
		if baseURL == "" || agentID == "" {
			return tokenanalysis.Client{}, true, errors.New("Agent Token analysis environment is incomplete")
		}
		return tokenanalysis.Client{
			BaseURL:    baseURL,
			AgentID:    agentID,
			AgentToken: agentToken,
		}, true, nil
	}
	status := daemonStatus()
	if status == nil || !status.Running || !status.Healthy || strings.TrimSpace(status.BaseURL) == "" {
		return tokenanalysis.Client{}, false, errors.New("daemon_unavailable: start AgentTeam Desktop or `agentteam daemon` first")
	}
	apiKey, err := token.Default().GenerateTempKey()
	if err != nil {
		return tokenanalysis.Client{}, false, fmt.Errorf("create local Token analysis credential: %w", err)
	}
	return tokenanalysis.Client{BaseURL: status.BaseURL, APIKey: apiKey}, false, nil
}

func writeTokenReport(path string, report tokenanalysis.Report) error {
	data, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		return err
	}
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
	if err != nil {
		return fmt.Errorf("write token report: %w", err)
	}
	if err := file.Chmod(0o600); err != nil {
		_ = file.Close()
		return err
	}
	if _, err := file.Write(append(data, '\n')); err != nil {
		_ = file.Close()
		return err
	}
	return file.Close()
}

func renderTokenReport(command *cobra.Command, report tokenanalysis.Report, format string) error {
	switch strings.ToLower(strings.TrimSpace(format)) {
	case "json":
		encoder := json.NewEncoder(command.OutOrStdout())
		encoder.SetIndent("", "  ")
		return encoder.Encode(report)
	case "table":
	default:
		return fmt.Errorf("unsupported format %q (valid: table|json)", format)
	}
	out := command.OutOrStdout()
	fmt.Fprintln(out, "AgentTeam Token 分析")
	fmt.Fprintf(out, "窗口: %s\n", report.Window)
	fmt.Fprintf(out, "快照: %s\n", report.SnapshotAt.Format(time.RFC3339))
	if report.Scope.Kind == "run" {
		fmt.Fprintf(out, "Run 状态: %s（快照后继续运行会产生新 usage；比较前请确认 ended）\n", report.Scope.RunState)
	}
	switch {
	case report.Collection.AutoEnabled:
		fmt.Fprintln(out, "详细采集: 已自动开启（不追溯历史；请复现问题后再次运行 analyze）")
	case report.Collection.Stopped:
		fmt.Fprintln(out, "详细采集: 已在本次报告后关闭")
	case report.Collection.Enabled:
		fmt.Fprintln(out, "详细采集: 开启")
	default:
		fmt.Fprintln(out, "详细采集: 关闭")
		if !report.Collection.CallerCanControl {
			fmt.Fprintln(out, "采集权限: 当前 Agent 只读；请由 Human 在本机 Token UI 开启，或从 Human shell 运行 `agentteam token analyze --window 1d`，再复现目标任务")
		}
	}
	writer := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
	fmt.Fprintf(writer, "Runtime 累计 input\t%s\n", formatTokenCount(report.Summary.RuntimeInput))
	fmt.Fprintf(writer, "Runtime cache read\t%s\n", formatTokenCount(report.Summary.RuntimeCacheRead))
	fmt.Fprintf(writer, "Runtime output\t%s\n", formatTokenCount(report.Summary.RuntimeOutput))
	fmt.Fprintf(writer, "Root 模型调用数\t%s\n", formatMetric(report.Summary.RootModelInvocations))
	fmt.Fprintf(writer, "Root 平均每次调用 input\t%s\n", formatMetric(report.Summary.RootAverageInputPerInvocation))
	fmt.Fprintf(writer, "Root runtime input\t%s\n", formatMetric(report.Summary.RootRuntimeInput))
	fmt.Fprintf(writer, "AgentTeam 可控 Context\t%s\n", formatMetric(report.Summary.ControllableMaterialized))
	fmt.Fprintf(writer, "AgentTeam 可控 exposure\t%s\n", formatMetric(report.Summary.ControllableExposure))
	fmt.Fprintf(writer, "Workspace 必要配置\t%s\n", formatMetric(report.Summary.ContextPartition.WorkspaceConfiguration))
	fmt.Fprintf(writer, "用户任务 payload\t%s\n", formatMetric(report.Summary.ContextPartition.UserTaskPayload))
	if err := writer.Flush(); err != nil {
		return err
	}
	if len(report.Opportunities) > 0 {
		fmt.Fprintln(out, "\n首要 AgentTeam 可控优化机会:")
		for index, opportunity := range report.Opportunities {
			fmt.Fprintf(out, "%d. %s — %s estimated (%s)\n   %s\n",
				index+1, opportunity.DisplayName, formatTokenCount(opportunity.Estimate), opportunity.Estimator, opportunity.SuggestedAction)
		}
	}
	if len(report.DataQuality.Diagnostics) > 0 {
		fmt.Fprintf(out, "\n数据质量: %s (%s)\n", report.DataQuality.State, strings.Join(report.DataQuality.Diagnostics, ", "))
	}
	return nil
}

func renderTokenComparison(command *cobra.Command, comparison tokenanalysis.Comparison, format string) error {
	switch strings.ToLower(strings.TrimSpace(format)) {
	case "json":
		encoder := json.NewEncoder(command.OutOrStdout())
		encoder.SetIndent("", "  ")
		return encoder.Encode(comparison)
	case "table":
	default:
		return fmt.Errorf("unsupported format %q (valid: table|json)", format)
	}
	fmt.Fprintf(command.OutOrStdout(), "AgentTeam Token 对比\n状态: %s\n", comparison.State)
	if len(comparison.MismatchReasons) > 0 {
		fmt.Fprintf(command.OutOrStdout(), "不匹配: %s\n", strings.Join(comparison.MismatchReasons, ", "))
		return nil
	}
	fmt.Fprintf(command.OutOrStdout(), "累计 input delta: %+d\n", comparison.InputDelta)
	if comparison.RootInvocationDelta != nil {
		fmt.Fprintf(command.OutOrStdout(), "Root 模型调用数 delta: %+d\n", *comparison.RootInvocationDelta)
	}
	fmt.Fprintln(command.OutOrStdout(), comparison.Note)
	return nil
}

func formatMetric(value tokenanalysis.MetricValue) string {
	if value.State != tokenanalysis.StateAvailable || value.Value == nil {
		return value.State
	}
	return formatTokenCount(*value.Value)
}

func formatTokenCount(value uint64) string {
	switch {
	case value >= 1_000_000:
		return fmt.Sprintf("%.1fM", float64(value)/1_000_000)
	case value >= 1_000:
		return fmt.Sprintf("%.1fK", float64(value)/1_000)
	default:
		return fmt.Sprint(value)
	}
}