package cmd

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

	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"github.com/tidbcloud/investibook/cli/internal/api"
)

var (
	dimColor  = color.New(color.FgHiBlack)
	blueColor = color.New(color.FgBlue)
	redColor  = color.New(color.FgRed)
)

var chatCmd = &cobra.Command{
	Use:   "chat",
	Short: "Chat with the AI agent",
	Long:  "Send a message to the AI agent and stream its response.\nUse --interactive for REPL mode.",
	Args:  cobra.MaximumNArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		interactive, _ := cmd.Flags().GetBool("interactive")
		detach, _ := cmd.Flags().GetBool("detach")
		client, cfg := mustClient()
		wsID := mustWorkspace(cfg)

		if interactive {
			runInteractiveChat(client, wsID)
			return
		}
		if len(args) == 0 {
			fmt.Fprintln(os.Stderr, "Provide a message or use --interactive")
			os.Exit(1)
		}
		if err := sendChatMessage(client, wsID, args[0], detach); err != nil {
			redColor.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
	},
}

func sendChatMessage(client *api.Client, wsID, message string, detach bool) error {
	userMessageID := fmt.Sprintf("cli-user-%d", time.Now().UnixMilli())
	run, err := client.CreateAgentRun(message, wsID, userMessageID)
	if err != nil {
		return err
	}
	if detach {
		fmt.Printf("Agent turn submitted: %s\n", run.ID)
		return nil
	}
	return streamAgentRun(client, run)
}

func streamAgentRun(client *api.Client, run *api.AgentRun) error {
	ctx, stopStream := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stopStream()
	ch, err := client.AgentRunEvents(ctx, run)
	if err != nil {
		if ctx.Err() != nil {
			return cancelAgentRun(client, run)
		}
		return err
	}

	toolStates := map[string]string{}
	var runError string
	for {
		var event api.SSEEvent
		var ok bool
		select {
		case <-ctx.Done():
			stopStream()
			return cancelAgentRun(client, run)
		case event, ok = <-ch:
			if !ok {
				if runError != "" {
					return errors.New(runError)
				}
				return nil
			}
		}
		switch event.Type {
		case "text":
			var d api.SSETextData
			if json.Unmarshal(event.Data, &d) == nil {
				fmt.Print(d.Text)
			}
		case "thought":
			var d api.SSETextData
			if json.Unmarshal(event.Data, &d) == nil {
				dimColor.Printf("[thinking] %s\n", d.Text)
			}
		case "tool":
			var d api.SSEToolData
			if json.Unmarshal(event.Data, &d) == nil {
				lastStatus := toolStates[d.ID]
				toolStates[d.ID] = d.Status
				if d.Status == "running" && lastStatus != "running" {
					title := d.Title
					if title == "" {
						title = d.ID
					}
					blueColor.Printf("[tool] %s", title)
					if d.Input != "" {
						blueColor.Printf("(%s)", truncate(d.Input, 120))
					}
					fmt.Println()
				} else if d.Status == "completed" && d.Output != "" {
					dimColor.Printf("  -> %s\n", truncate(d.Output, 200))
				}
			}
		case "error":
			var d api.SSEErrorData
			if json.Unmarshal(event.Data, &d) == nil {
				runError = d.Error
			}
		case "complete":
			fmt.Println()
			var d api.SSECompleteData
			if json.Unmarshal(event.Data, &d) == nil && (d.Status == "failed" || d.Status == "interrupted") {
				if runError == "" {
					runError = "agent turn " + d.Status
				}
				return errors.New(runError)
			}
			return nil
		case "session_created":
			// silently consume
		}
	}
}

func cancelAgentRun(client *api.Client, run *api.AgentRun) error {
	if _, err := client.CancelAgentRun(run.SessionID); err != nil {
		return fmt.Errorf("cancel turn %s: %w", run.ID, err)
	}
	dimColor.Fprintf(os.Stderr, "Cancelled agent turn %s\n", run.ID)
	return nil
}

func runInteractiveChat(client *api.Client, wsID string) {
	scanner := bufio.NewScanner(os.Stdin)
	fmt.Println("Investibook Chat (Ctrl+D or /exit to quit)")
	fmt.Println()
	for {
		fmt.Print("> ")
		if !scanner.Scan() {
			break
		}
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		if line == "/exit" || line == "/quit" {
			break
		}

		run, err := client.CreateAgentRun(
			line,
			wsID,
			fmt.Sprintf("cli-user-%d", time.Now().UnixMilli()),
		)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			fmt.Println()
			continue
		}
		if err := streamAgentRun(client, run); err != nil {
			redColor.Fprintf(os.Stderr, "Error: %v\n", err)
		}
		fmt.Println()
	}
	fmt.Println("\nGoodbye!")
}

func truncate(s string, max int) string {
	s = strings.ReplaceAll(s, "\n", " ")
	if len(s) > max {
		return s[:max] + "..."
	}
	return s
}

func init() {
	chatCmd.Flags().BoolP("interactive", "i", false, "start interactive REPL mode")
	chatCmd.Flags().Bool("detach", false, "submit the turn and return without streaming")
	rootCmd.AddCommand(chatCmd)
}