package agent

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"os"
	"strings"
	"time"

	"github.com/henrygd/beszel"
	"github.com/henrygd/beszel/agent/utils"
	"github.com/henrygd/beszel/internal/common"
	"github.com/henrygd/beszel/internal/entities/system"

	"github.com/blang/semver"
	"github.com/fxamacker/cbor/v2"
	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
)

// ServerOptions contains configuration options for starting the SSH server.
type ServerOptions struct {
	Addr    string            // Network address to listen on (e.g., ":45876" or "/path/to/socket")
	Network string            // Network type ("tcp" or "unix")
	Keys    []gossh.PublicKey // SSH public keys for authentication
}

// hubVersions caches hub versions by session ID to avoid repeated parsing.
var hubVersions map[string]semver.Version

// StartServer starts the SSH server with the provided options.
// It configures the server with secure defaults, sets up authentication,
// and begins listening for connections. Returns an error if the server
// is already running or if there's an issue starting the server.
func (a *Agent) StartServer(opts ServerOptions) error {
	if disableSSH, _ := utils.GetEnv("DISABLE_SSH"); disableSSH == "true" {
		return errors.New("SSH disabled")
	}
	if a.server != nil {
		return errors.New("server already started")
	}

	slog.Info("Starting SSH server", "addr", opts.Addr, "network", opts.Network)

	if opts.Network == "unix" {
		// remove existing socket file if it exists
		if err := os.Remove(opts.Addr); err != nil && !os.IsNotExist(err) {
			return err
		}
	}

	// start listening on the address
	ln, err := net.Listen(opts.Network, opts.Addr)
	if err != nil {
		return err
	}
	defer ln.Close()

	// base config (limit to allowed algorithms)
	config := &gossh.ServerConfig{
		ServerVersion: fmt.Sprintf("SSH-2.0-%s_%s", beszel.AppName, beszel.Version),
	}
	config.KeyExchanges = common.DefaultKeyExchanges
	config.MACs = common.DefaultMACs
	config.Ciphers = common.DefaultCiphers

	// set default handler
	ssh.Handle(a.handleSession)

	a.server = &ssh.Server{
		ServerConfigCallback: func(ctx ssh.Context) *gossh.ServerConfig {
			return config
		},
		// check public key(s)
		PublicKeyHandler: func(ctx ssh.Context, key ssh.PublicKey) bool {
			remoteAddr := ctx.RemoteAddr()
			for _, pubKey := range opts.Keys {
				if ssh.KeysEqual(key, pubKey) {
					slog.Info("SSH connected", "addr", remoteAddr)
					return true
				}
			}
			slog.Warn("Invalid SSH key", "addr", remoteAddr)
			return false
		},
		// disable pty
		PtyCallback: func(ctx ssh.Context, pty ssh.Pty) bool {
			return false
		},
		// close idle connections after 70 seconds
		IdleTimeout: 70 * time.Second,
	}

	// Start SSH server on the listener
	return a.server.Serve(ln)
}

// getHubVersion retrieves and caches the hub version for a given session.
// It extracts the version from the SSH client version string and caches
// it to avoid repeated parsing. Returns a zero version if parsing fails.
func (a *Agent) getHubVersion(sessionId string, sessionCtx ssh.Context) semver.Version {
	if hubVersions == nil {
		hubVersions = make(map[string]semver.Version, 1)
	}
	hubVersion, ok := hubVersions[sessionId]
	if ok {
		return hubVersion
	}
	// Extract hub version from SSH client version
	clientVersion := sessionCtx.Value(ssh.ContextKeyClientVersion)
	if versionStr, ok := clientVersion.(string); ok {
		hubVersion, _ = extractHubVersion(versionStr)
	}
	hubVersions[sessionId] = hubVersion
	return hubVersion
}

// handleSession handles an incoming SSH session by gathering system statistics
// and sending them to the hub. It signals connection events, determines the
// appropriate encoding format based on hub version, and exits with appropriate
// status codes.
func (a *Agent) handleSession(s ssh.Session) {
	a.connectionManager.eventChan <- SSHConnect

	sessionCtx := s.Context()
	sessionID := sessionCtx.SessionID()

	hubVersion := a.getHubVersion(sessionID, sessionCtx)

	// Legacy one-shot behavior for older hubs
	if hubVersion.LT(beszel.MinVersionAgentResponse) {
		if err := a.handleLegacyStats(s, hubVersion); err != nil {
			slog.Error("Error encoding stats", "err", err)
			s.Exit(1)
			return
		}
	}

	var req common.HubRequest[cbor.RawMessage]
	if err := cbor.NewDecoder(s).Decode(&req); err != nil {
		// Fallback to legacy one-shot if the first decode fails
		if err2 := a.handleLegacyStats(s, hubVersion); err2 != nil {
			slog.Error("Error encoding stats (fallback)", "err", err2)
			s.Exit(1)
			return
		}
		s.Exit(0)
		return
	}
	if err := a.handleSSHRequest(s, &req); err != nil {
		slog.Error("SSH request handling failed", "err", err)
		s.Exit(1)
		return
	}
	s.Exit(0)
}

// handleSSHRequest builds a handler context and dispatches to the shared registry
func (a *Agent) handleSSHRequest(w io.Writer, req *common.HubRequest[cbor.RawMessage]) error {
	// SSH does not support fingerprint auth action
	if req.Action == common.CheckFingerprint {
		return cbor.NewEncoder(w).Encode(common.AgentResponse{Error: "unsupported action"})
	}

	// responder that writes AgentResponse to stdout
	// Uses legacy typed fields for backward compatibility with <= 0.17
	sshResponder := func(data any, requestID *uint32) error {
		response := newAgentResponse(data, requestID)
		return cbor.NewEncoder(w).Encode(response)
	}

	ctx := &HandlerContext{
		Client:       nil,
		Agent:        a,
		Request:      req,
		RequestID:    nil,
		HubVerified:  true,
		SendResponse: sshResponder,
	}

	if handler, ok := a.handlerRegistry.GetHandler(req.Action); ok {
		if err := handler.Handle(ctx); err != nil {
			return cbor.NewEncoder(w).Encode(common.AgentResponse{Error: err.Error()})
		}
		return nil
	}
	return cbor.NewEncoder(w).Encode(common.AgentResponse{Error: fmt.Sprintf("unknown action: %d", req.Action)})
}

// handleLegacyStats serves the legacy one-shot stats payload for older hubs
func (a *Agent) handleLegacyStats(w io.Writer, hubVersion semver.Version) error {
	stats := a.gatherStats(common.DataRequestOptions{CacheTimeMs: defaultDataCacheTimeMs})
	return a.writeToSession(w, stats, hubVersion)
}

// writeToSession encodes and writes system statistics to the session.
// It chooses between CBOR and JSON encoding based on the hub version,
// using CBOR for newer versions and JSON for legacy compatibility.
func (a *Agent) writeToSession(w io.Writer, stats *system.CombinedData, hubVersion semver.Version) error {
	if hubVersion.GTE(beszel.MinVersionCbor) {
		return cbor.NewEncoder(w).Encode(stats)
	}
	return json.NewEncoder(w).Encode(stats)
}

// extractHubVersion extracts the beszel version from SSH client version string.
// Expected format: "SSH-2.0-beszel_X.Y.Z" or "beszel_X.Y.Z"
func extractHubVersion(versionString string) (semver.Version, error) {
	_, after, _ := strings.Cut(versionString, "_")
	return semver.Parse(after)
}

// ParseKeys parses a string containing SSH public keys in authorized_keys format.
// It returns a slice of ssh.PublicKey and an error if any key fails to parse.
func ParseKeys(input string) ([]gossh.PublicKey, error) {
	var parsedKeys []gossh.PublicKey
	for line := range strings.Lines(input) {
		line = strings.TrimSpace(line)
		// Skip empty lines or comments
		if len(line) == 0 || strings.HasPrefix(line, "#") {
			continue
		}
		// Parse the key
		parsedKey, _, _, _, err := gossh.ParseAuthorizedKey([]byte(line))
		if err != nil {
			return nil, fmt.Errorf("failed to parse key: %s, error: %w", line, err)
		}
		parsedKeys = append(parsedKeys, parsedKey)
	}
	return parsedKeys, nil
}

// GetAddress determines the network address to listen on from various sources.
// It checks the provided address, then environment variables (LISTEN, PORT),
// and finally defaults to ":45876".
func GetAddress(addr string) string {
	if addr == "" {
		addr, _ = utils.GetEnv("LISTEN")
	}
	if addr == "" {
		// Legacy PORT environment variable support
		addr, _ = utils.GetEnv("PORT")
	}
	if addr == "" {
		return ":45876"
	}
	// prefix with : if only port was provided
	if GetNetwork(addr) != "unix" && !strings.Contains(addr, ":") {
		addr = ":" + addr
	}
	return addr
}

// GetNetwork determines the network type based on the address format.
// It checks the NETWORK environment variable first, then infers from
// the address format: addresses starting with "/" are "unix", others are "tcp".
func GetNetwork(addr string) string {
	if network, ok := utils.GetEnv("NETWORK"); ok && network != "" {
		return network
	}
	if strings.HasPrefix(addr, "/") {
		return "unix"
	}
	return "tcp"
}

// StopServer stops the SSH server if it's running.
// It returns an error if the server is not running or if there's an error stopping it.
func (a *Agent) StopServer() error {
	if a.server == nil {
		return errors.New("SSH server not running")
	}

	slog.Info("Stopping SSH server")
	_ = a.server.Close()
	a.server = nil
	a.connectionManager.eventChan <- SSHDisconnect
	return nil
}


import * as React from "react"

import { cn } from "#/lib/utils.ts"
import { HugeiconsIcon } from "@hugeicons/react"
import { UnfoldMoreIcon } from "@hugeicons/core-free-icons"

type NativeSelectProps = Omit<React.ComponentProps<"select">, "size"> & {
  size?: "sm" | "default"
}

function NativeSelect({
  className,
  size = "default",
  ...props
}: NativeSelectProps) {
  return (
    <div
      className={cn(
        "group/native-select relative w-fit has-[select:disabled]:opacity-50",
        className
      )}
      data-slot="native-select-wrapper"
      data-size={size}
    >
      <select
        data-slot="native-select"
        data-size={size}
        className="h-7 w-full min-w-0 appearance-none rounded-md border border-input bg-input/20 py-0.5 pr-6 pl-2 text-xs/relaxed transition-colors outline-none select-none selection:bg-primary selection:text-primary-foreground placeholder:text-muted-foreground focus-visible:border-ring focus-visible:ring-2 focus-visible:ring-ring/30 disabled:pointer-events-none disabled:cursor-not-allowed aria-invalid:border-destructive aria-invalid:ring-2 aria-invalid:ring-destructive/20 data-[size=sm]:h-6 data-[size=sm]:text-[0.625rem] dark:bg-input/30 dark:hover:bg-input/50 dark:aria-invalid:border-destructive/50 dark:aria-invalid:ring-destructive/40"
        {...props}
      />
      <HugeiconsIcon icon={UnfoldMoreIcon} strokeWidth={2} className="pointer-events-none absolute top-1/2 right-1.5 size-3.5 -translate-y-1/2 text-muted-foreground select-none group-data-[size=sm]/native-select:size-3 group-data-[size=sm]/native-select:-translate-y-[calc(--spacing(1.25))]" aria-hidden="true" data-slot="native-select-icon" />
    </div>
  )
}

function NativeSelectOption({
  className,
  ...props
}: React.ComponentProps<"option">) {
  return (
    <option
      data-slot="native-select-option"
      className={cn("bg-[Canvas] text-[CanvasText]", className)}
      {...props}
    />
  )
}

function NativeSelectOptGroup({
  className,
  ...props
}: React.ComponentProps<"optgroup">) {
  return (
    <optgroup
      data-slot="native-select-optgroup"
      className={cn("bg-[Canvas] text-[CanvasText]", className)}
      {...props}
    />
  )
}

export { NativeSelect, NativeSelectOptGroup, NativeSelectOption }


"""
ChatLiteLLM - LiteLLM chat model wrapper.

Requires the `litellm` package to be installed separately:
    pip install litellm

Note: litellm is NOT included as a dependency of browser-use.
"""

import logging
from dataclasses import dataclass, field
from typing import Any, TypeVar, overload

from pydantic import BaseModel

from browser_use.llm.base import BaseChatModel
from browser_use.llm.exceptions import ModelProviderError, ModelRateLimitError
from browser_use.llm.messages import BaseMessage
from browser_use.llm.schema import SchemaOptimizer
from browser_use.llm.views import ChatInvokeCompletion, ChatInvokeUsage

from .serializer import LiteLLMMessageSerializer

logger = logging.getLogger(__name__)

T = TypeVar('T', bound=BaseModel)

@dataclass
class ChatLiteLLM(BaseChatModel):
	model: str
	api_key: str | None = None
	api_base: str | None = None
	temperature: float | None = 0.0
	max_tokens: int | None = 4096
	max_retries: int = 3
	metadata: dict[str, Any] | None = None

	_provider_name: str = field(default='', init=False, repr=False)
	_clean_model: str = field(default='', init=False, repr=False)

	def __post_init__(self) -> None:
		"""Resolve provider info from the model string via litellm."""
		try:
			from litellm import get_llm_provider  # type: ignore[reportMissingImports]

			self._clean_model, self._provider_name, _, _ = get_llm_provider(self.model)
		except Exception:
			if '/' in self.model:
				self._provider_name, self._clean_model = self.model.split('/', 1)
			else:
				self._provider_name = 'openai'
				self._clean_model = self.model

		logger.debug(
			'ChatLiteLLM initialized: model=%s, provider=%s, clean=%s, api_base=%s',
			self.model,
			self._provider_name,
			self._clean_model,
			self.api_base or '(default)',
		)

	@property
	def provider(self) -> str:
		return self._provider_name or 'litellm'

	@property
	def name(self) -> str:
		return self._clean_model or self.model

	@staticmethod
	def _parse_usage(response: Any) -> ChatInvokeUsage | None:
		"""Extract token usage from a litellm response."""
		usage = getattr(response, 'usage', None)
		if usage is None:
			return None

		prompt_tokens = getattr(usage, 'prompt_tokens', 0) or 0
		completion_tokens = getattr(usage, 'completion_tokens', 0) or 0

		prompt_cached = getattr(usage, 'cache_read_input_tokens', None)
		cache_creation = getattr(usage, 'cache_creation_input_tokens', None)

		if prompt_cached is None:
			details = getattr(usage, 'prompt_tokens_details', None)
			if details:
				prompt_cached = getattr(details, 'cached_tokens', None)

		return ChatInvokeUsage(
			prompt_tokens=prompt_tokens,
			prompt_cached_tokens=int(prompt_cached) if prompt_cached is not None else None,
			prompt_cache_creation_tokens=int(cache_creation) if cache_creation is not None else None,
			prompt_image_tokens=None,
			completion_tokens=completion_tokens,
			total_tokens=prompt_tokens + completion_tokens,
		)

	@overload
	async def ainvoke(
		self,
		messages: list[BaseMessage],
		output_format: None = None,
		**kwargs: Any,
	) -> ChatInvokeCompletion[str]: ...

	@overload
	async def ainvoke(
		self,
		messages: list[BaseMessage],
		output_format: type[T],
		**kwargs: Any,
	) -> ChatInvokeCompletion[T]: ...

	async def ainvoke(
		self,
		messages: list[BaseMessage],
		output_format: type[T] | None = None,
		**kwargs: Any,
	) -> ChatInvokeCompletion[T] | ChatInvokeCompletion[str]:
		from litellm import acompletion  # type: ignore[reportMissingImports]
		from litellm.exceptions import APIConnectionError, APIError, RateLimitError, Timeout  # type: ignore[reportMissingImports]
		from litellm.types.utils import ModelResponse  # type: ignore[reportMissingImports]

		litellm_messages = LiteLLMMessageSerializer.serialize(messages)

		params: dict[str, Any] = {
			'model': self.model,
			'messages': litellm_messages,
			'num_retries': self.max_retries,
		}

		if self.temperature is not None:
			params['temperature'] = self.temperature
		if self.max_tokens is not None:
			params['max_tokens'] = self.max_tokens
		if self.api_key:
			params['api_key'] = self.api_key
		if self.api_base:
			params['api_base'] = self.api_base
		if self.metadata:
			params['metadata'] = self.metadata

		if output_format is not None:
			schema = SchemaOptimizer.create_optimized_json_schema(output_format)
			params['response_format'] = {
				'type': 'json_schema',
				'json_schema': {
					'name': 'agent_output',
					'strict': True,
					'schema': schema,
				},
			}

		try:
			raw_response = await acompletion(**params)
		except RateLimitError as e:
			raise ModelRateLimitError(
				message=str(e),
				model=self.name,
			) from e
		except Timeout as e:
			raise ModelProviderError(
				message=f'Request timed out: {e}',
				model=self.name,
			) from e
		except APIConnectionError as e:
			raise ModelProviderError(
				message=str(e),
				model=self.name,
			) from e
		except APIError as e:
			status = getattr(e, 'status_code', 502) or 502
			raise ModelProviderError(
				message=str(e),
				status_code=status,
				model=self.name,
			) from e
		except ModelProviderError:
			raise
		except Exception as e:
			raise ModelProviderError(
				message=str(e),
				model=self.name,
			) from e

		assert isinstance(raw_response, ModelResponse), f'Expected ModelResponse, got {type(raw_response)}'
		response: ModelResponse = raw_response

		choice = response.choices[0] if response.choices else None
		if choice is None:
			raise ModelProviderError(
				message='Empty response: no choices returned by the model',
				status_code=502,
				model=self.name,
			)

		content = choice.message.content or ''
		usage = self._parse_usage(response)
		stop_reason = choice.finish_reason

		thinking: str | None = None
		msg_obj = choice.message
		reasoning = getattr(msg_obj, 'reasoning_content', None)
		if reasoning:
			thinking = str(reasoning)

		if output_format is not None:
			if not content:
				raise ModelProviderError(
					message='Model returned empty content for structured output request',
					status_code=500,
					model=self.name,
				)
			parsed = output_format.model_validate_json(content)
			return ChatInvokeCompletion(
				completion=parsed,
				thinking=thinking,
				usage=usage,
				stop_reason=stop_reason,
			)

		return ChatInvokeCompletion(
			completion=content,
			thinking=thinking,
			usage=usage,
			stop_reason=stop_reason,
		)


#!/usr/bin/env node
// Tests for src/tools/caveman-init.js — fixture-based.
// Run: node tests/test_caveman_init.js

const fs = require('fs');
const path = require('path');
const os = require('os');
const assert = require('assert');
const { execFileSync } = require('child_process');

const ROOT = path.resolve(__dirname, '..');
const INIT = path.join(ROOT, 'src', 'tools', 'caveman-init.js');

let passed = 0;
let failed = 0;

// Point OPENCLAW_WORKSPACE at a nonexistent dir inside the fixture so the
// openclaw target reports skipped-workspace-missing instead of writing to
// the developer's real ~/.openclaw/workspace.
function runInit(tmp, ...args) {
  return execFileSync(process.execPath, [INIT, tmp, ...args], {
    encoding: 'utf8',
    env: { ...process.env, OPENCLAW_WORKSPACE: path.join(tmp, 'no-openclaw') },
  });
}

function test(name, fn) {
  const tmp = fs.mkdtempSync(path.join(os.tmpdir(), 'caveman-init-test-'));
  try {
    fn(tmp);
    passed++;
    console.log(`  ✓ ${name}`);
  } catch (e) {
    failed++;
    console.error(`  ✗ ${name}\n    ${e.message}`);
  } finally {
    fs.rmSync(tmp, { recursive: true, force: true });
  }
}

console.log('caveman-init tests\n');

test('greenfield: creates all rule files with proper frontmatter', (tmp) => {
  runInit(tmp);
  const cursor = fs.readFileSync(path.join(tmp, '.cursor/rules/caveman.mdc'), 'utf8');
  assert.match(cursor, /alwaysApply: true/);
  assert.match(cursor, /Respond terse like smart caveman/);
  const windsurf = fs.readFileSync(path.join(tmp, '.windsurf/rules/caveman.md'), 'utf8');
  assert.match(windsurf, /trigger: always_on/);
  const cline = fs.readFileSync(path.join(tmp, '.clinerules/caveman.md'), 'utf8');
  assert.match(cline, /^Respond terse/);
  const copilot = fs.readFileSync(path.join(tmp, '.github/copilot-instructions.md'), 'utf8');
  assert.match(copilot, /Respond terse/);
  const agents = fs.readFileSync(path.join(tmp, 'AGENTS.md'), 'utf8');
  assert.match(agents, /Respond terse/);
  const opencode = fs.readFileSync(path.join(tmp, '.opencode/AGENTS.md'), 'utf8');
  assert.match(opencode, /Respond terse/);
});

test('idempotent: re-running on a clean install skips all', (tmp) => {
  runInit(tmp);
  const out = runInit(tmp);
  // 6 repo rule files skipped-already-installed + openclaw skipped (no workspace)
  assert.match(out, /7 skipped/);
  assert.doesNotMatch(out, /[1-9]\d* added/);
});

test('append mode: existing AGENTS.md gets caveman appended (not replaced)', (tmp) => {
  fs.writeFileSync(path.join(tmp, 'AGENTS.md'), '# My project\n\nDo not delete me.\n');
  runInit(tmp);
  const agents = fs.readFileSync(path.join(tmp, 'AGENTS.md'), 'utf8');
  assert.match(agents, /Do not delete me/);
  assert.match(agents, /Respond terse like smart caveman/);
});

test('skip mode: existing .cursor rule is not overwritten without --force', (tmp) => {
  const dir = path.join(tmp, '.cursor/rules');
  fs.mkdirSync(dir, { recursive: true });
  fs.writeFileSync(path.join(dir, 'caveman.mdc'), '# original\nDo not delete me.\n');
  const out = runInit(tmp);
  assert.match(out, /\? .*\.cursor\/rules\/caveman\.mdc/);
  const after = fs.readFileSync(path.join(dir, 'caveman.mdc'), 'utf8');
  assert.strictEqual(after, '# original\nDo not delete me.\n');
});

test('--force overwrites existing rule files', (tmp) => {
  const dir = path.join(tmp, '.cursor/rules');
  fs.mkdirSync(dir, { recursive: true });
  fs.writeFileSync(path.join(dir, 'caveman.mdc'), '# original\n');
  runInit(tmp, '--force');
  const after = fs.readFileSync(path.join(dir, 'caveman.mdc'), 'utf8');
  assert.match(after, /alwaysApply: true/);
  assert.match(after, /Respond terse/);
});

test('--dry-run: announces but writes nothing', (tmp) => {
  const out = runInit(tmp, '--dry-run');
  assert.match(out, /\(dry run\)/);
  assert.match(out, /6 added/);
  assert.ok(!fs.existsSync(path.join(tmp, '.cursor')));
  assert.ok(!fs.existsSync(path.join(tmp, '.windsurf')));
  assert.ok(!fs.existsSync(path.join(tmp, '.clinerules')));
  assert.ok(!fs.existsSync(path.join(tmp, '.github/copilot-instructions.md')));
  assert.ok(!fs.existsSync(path.join(tmp, '.opencode')));
  assert.ok(!fs.existsSync(path.join(tmp, 'AGENTS.md')));
});

test('--only filters to one target', (tmp) => {
  const out = runInit(tmp, '--only', 'cline');
  assert.match(out, /1 added/);
  assert.ok(fs.existsSync(path.join(tmp, '.clinerules/caveman.md')));
  assert.ok(!fs.existsSync(path.join(tmp, '.cursor')));
});

test('detects sentinel and skips files that already have caveman content', (tmp) => {
  // Hand-write a file that already contains the rule (simulating prior install).
  const dir = path.join(tmp, '.clinerules');
  fs.mkdirSync(dir, { recursive: true });
  fs.writeFileSync(path.join(dir, 'caveman.md'),
    '# Existing\n\nRespond terse like smart caveman. Hello.\n');
  const out = runInit(tmp, '--only', 'cline');
  assert.match(out, /skipped-already-installed/);
});

console.log(`\n${passed} passed, ${failed} failed`);
process.exit(failed ? 1 : 0);


//! Conversation memory: Rig-managed persistent conversation history for agents.
//!
//! Memory differs from existing agent context features:
//! - classic runtime context: static documents always included in prompts;
//! - classic runtime request patches: per-turn documents supplied by application hooks;
//! - caller-managed message history supplied directly on completion requests;
//! - **Memory** (this module): Rig-managed history loaded and saved automatically per
//!   conversation id.
//!
//! # Example
//!
//! ```no_run
//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
//! use rig_core::{
//!     completion::Message,
//!     memory::{ConversationMemory, InMemoryConversationMemory},
//! };
//!
//! let memory = InMemoryConversationMemory::new();
//! memory
//!     .append(
//!         "thread-1",
//!         vec![
//!             Message::user("My name is Alice."),
//!             Message::assistant("Hello, Alice!"),
//!         ],
//!     )
//!     .await?;
//! let history = memory.load("thread-1").await?;
//! assert_eq!(history.len(), 2);
//! # Ok(()) }
//! ```
//!
//! Truncation, summarization, and other history-shaping policies live in the
//! `rig-memory` companion crate. To shape history inside the in-tree backend,
//! pass a closure to [`InMemoryConversationMemory::with_filter`].

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

use crate::{
    completion::Message,
    wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync},
};

/// Boxed error source for memory backend failures.
#[cfg(not(target_family = "wasm"))]
pub type MemoryBackendError = Box<dyn std::error::Error + Send + Sync + 'static>;

/// Boxed error source for memory backend failures.
#[cfg(target_family = "wasm")]
pub type MemoryBackendError = Box<dyn std::error::Error + 'static>;

/// Errors produced by a [`ConversationMemory`] backend.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum MemoryError {
    /// The backing store failed to load, append, or clear messages.
    #[error("Memory backend error: {0}")]
    Backend(MemoryBackendError),

    /// A history-shaping filter or policy rejected the loaded history.
    #[error("Memory policy error: {0}")]
    Policy(String),

    /// An internal invariant was violated (e.g. a poisoned in-process lock).
    /// Distinct from [`MemoryError::Backend`], which is reserved for failures
    /// of the underlying conversation store.
    #[error("Memory internal error: {0}")]
    Internal(String),
}

impl MemoryError {
    /// Wrap an arbitrary error from a backend implementation.
    pub fn backend<E>(source: E) -> Self
    where
        E: Into<MemoryBackendError>,
    {
        Self::Backend(source.into())
    }
}

/// A persistent conversation history backend.
///
/// Implementors store an ordered list of [`Message`]s per `conversation_id`. Rig
/// runtimes invoke [`ConversationMemory::load`] before sending a prompt and
/// [`ConversationMemory::append`] after a successful turn.
///
/// Implementations should keep `append` cheap; it runs inline before the agent
/// returns its response.
pub trait ConversationMemory: WasmCompatSend + WasmCompatSync {
    /// Load the full conversation history for `conversation_id`.
    ///
    /// Returns an empty `Vec` if the conversation has no stored messages.
    fn load<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>>;

    /// Append `messages` to the conversation identified by `conversation_id`.
    ///
    /// Called after a successful agent turn with the user prompt, the assistant
    /// response, and any tool-call/tool-result pairs that occurred during the turn.
    fn append<'a>(
        &'a self,
        conversation_id: &'a str,
        messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>>;

    /// Remove all stored messages for `conversation_id`.
    fn clear<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>>;
}

impl<M> ConversationMemory for Arc<M>
where
    M: ConversationMemory + ?Sized,
{
    fn load<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
        (**self).load(conversation_id)
    }

    fn append<'a>(
        &'a self,
        conversation_id: &'a str,
        messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        (**self).append(conversation_id, messages)
    }

    fn clear<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        (**self).clear(conversation_id)
    }
}

impl<M> ConversationMemory for Box<M>
where
    M: ConversationMemory + ?Sized,
{
    fn load<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
        (**self).load(conversation_id)
    }

    fn append<'a>(
        &'a self,
        conversation_id: &'a str,
        messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        (**self).append(conversation_id, messages)
    }

    fn clear<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        (**self).clear(conversation_id)
    }
}

/// A history-shaping closure applied during [`InMemoryConversationMemory::load`].
///
/// Implemented automatically for any closure with the right signature; the
/// trait exists to combine `Fn` with the WASM-compatible `Send`/`Sync` markers
/// in a single trait object.
pub trait MessageFilter:
    Fn(Vec<Message>) -> Vec<Message> + WasmCompatSend + WasmCompatSync
{
}

impl<F> MessageFilter for F where
    F: Fn(Vec<Message>) -> Vec<Message> + WasmCompatSend + WasmCompatSync
{
}

/// A side-channel for messages that a memory policy or adapter removes from
/// active history during [`ConversationMemory::load`].
///
/// Truncating policies (sliding window, token budget, …) drop older turns
/// once their limit is exceeded. Without a hook those messages are silently
/// lost. A [`DemotionHook`] receives the demoted messages and can persist
/// them into a long-tail store (semantic memory, episodic recall, archival
/// storage, …), turning truncation into demotion.
///
/// The trait is defined here in `rig-core` so that *any* memory backend
/// (in-memory, vector store, file archive, …) can implement it without
/// taking on a `rig-memory` dependency. The composing adapter that actually
/// wires a [`ConversationMemory`] backend, a policy, and a hook together
/// lives in the `rig-memory` companion crate.
///
/// Hooks should be inexpensive: their future is awaited inline on every
/// `load` that produces demoted messages, so a slow hook delays the agent's
/// next turn. Offload heavy I/O (network writes, disk fsyncs, …) to a
/// background task or a buffered channel inside the implementation.
///
/// # Idempotency contract
///
/// Implementations **must** be idempotent on the
/// `(conversation_id, messages)` pair. Composing adapters such as the
/// `DemotingPolicyMemory` in `rig-memory` track in-process delivery
/// watermarks to avoid replaying the same demotion within a single
/// process lifetime, but those watermarks are not persisted: across
/// process restarts (or when a new adapter is constructed over an
/// existing backend) the hook will receive previously-delivered
/// messages again. Hooks that append to durable storage should
/// deduplicate by content hash, by `(conversation_id, message_id)`,
/// or by an equivalent stable key.
pub trait DemotionHook: WasmCompatSend + WasmCompatSync {
    /// Receive `messages` that were demoted out of the active window for
    /// `conversation_id`.
    ///
    /// `messages` are in original conversation order. Errors are propagated
    /// as [`MemoryError::Backend`] by the composing adapter.
    fn on_demote<'a>(
        &'a self,
        conversation_id: &'a str,
        messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>>;
}

/// A [`DemotionHook`] that does nothing. Useful as a default when an adapter
/// requires a hook value but the caller has no long-tail store wired up yet.
#[derive(Debug, Default, Clone, Copy)]
pub struct NoopDemotionHook;

impl DemotionHook for NoopDemotionHook {
    fn on_demote<'a>(
        &'a self,
        _conversation_id: &'a str,
        _messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        Box::pin(async move { Ok(()) })
    }
}

/// Forwarding impl so callers can pass `Arc<H>` wherever a `DemotionHook`
/// is expected (e.g. when sharing a single hook between multiple memory
/// adapters).
impl<H> DemotionHook for Arc<H>
where
    H: DemotionHook + ?Sized,
{
    fn on_demote<'a>(
        &'a self,
        conversation_id: &'a str,
        messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        (**self).on_demote(conversation_id, messages)
    }
}

/// Derives a single [`Message`]-shaped artifact from a slice of messages
/// that a memory policy has evicted from the active window.
///
/// Where a [`DemotionHook`] is a one-way drain — observe what fell out and
/// return `()` — a `Compactor` is the inverse: it takes the evicted prefix
/// (and optionally the previous summary) and produces a derived artifact
/// that the composing adapter splices *back into* the active history. The
/// resulting prompt is no longer a verbatim suffix of the conversation; it
/// is `[summary, ...recent_window]`.
///
/// Implementations typically wrap an LLM call (`LlmCompactor<M>`) or a
/// pure template rollup. They run inline on the load path whenever the
/// policy demotes new messages, so a slow compactor delays the agent's
/// next turn — keep them fast or offload to a cached/background pipeline.
///
/// # Rolling summaries
///
/// `carry_over` is the artifact produced by the previous compaction for
/// this conversation, if any. Implementations that want a *recursive*
/// summary (the canonical pattern for long-running agents) should
/// summarize `evicted` *together with* `carry_over` so context lost in
/// earlier compactions is preserved transitively. Stateless implementations
/// can ignore `carry_over` and produce a fresh summary of `evicted` alone.
///
/// # Idempotency contract
///
/// Composing adapters track per-conversation in-process delivery so the
/// same `evicted` slice is not compacted twice within a process lifetime,
/// but those watermarks are not persisted across restarts. Implementations
/// that have side effects (writing summaries to a vector store, billing an
/// LLM call) should deduplicate by conversation id and content hash, the
/// same way [`DemotionHook`] implementations do.
pub trait Compactor: WasmCompatSend + WasmCompatSync {
    /// The summary value produced by [`Compactor::compact`].
    ///
    /// `Into<Message>` is required so the composing adapter can splice the
    /// artifact at the front of the loaded history. `Clone` is required so
    /// the adapter can keep a private copy as `carry_over` for the next
    /// compaction.
    type Artifact: Into<Message> + Clone + WasmCompatSend + WasmCompatSync + 'static;

    /// Produce a summary artifact for `evicted`, optionally combining it
    /// with the previous summary in `carry_over`.
    ///
    /// `evicted` is in original conversation order. Errors are propagated
    /// unchanged by composing adapters; pick the [`MemoryError`] variant
    /// that best describes the failure ([`MemoryError::Backend`] for I/O
    /// or remote-LLM faults, [`MemoryError::Internal`] for invariant
    /// breaks, and so on). The adapter does not re-wrap the returned
    /// variant.
    fn compact<'a>(
        &'a self,
        conversation_id: &'a str,
        evicted: &'a [Message],
        carry_over: Option<&'a Self::Artifact>,
    ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>>;
}

/// Forwarding impl so callers can pass `Arc<C>` wherever a `Compactor` is
/// expected (e.g. when sharing a single compactor across adapters).
impl<C> Compactor for Arc<C>
where
    C: Compactor + ?Sized,
{
    type Artifact = C::Artifact;

    fn compact<'a>(
        &'a self,
        conversation_id: &'a str,
        evicted: &'a [Message],
        carry_over: Option<&'a Self::Artifact>,
    ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
        (**self).compact(conversation_id, evicted, carry_over)
    }
}

/// A simple thread-safe in-memory [`ConversationMemory`] backed by a `HashMap`.
///
/// Messages are stored in process memory only and lost on restart. Useful for
/// tests, examples, and short-lived agents. Pass a closure to
/// [`InMemoryConversationMemory::with_filter`] to apply a history-shaping
/// transformation on every load (truncation, summarization, re-ordering, etc.).
/// Reusable named policies live in the `rig-memory` companion crate.
#[derive(Clone, Default)]
pub struct InMemoryConversationMemory {
    inner: Arc<Mutex<HashMap<String, Vec<Message>>>>,
    filter: Option<Arc<dyn MessageFilter>>,
}

impl InMemoryConversationMemory {
    /// Create an empty in-memory store with no filter.
    pub fn new() -> Self {
        Self::default()
    }

    /// Apply `filter` to the loaded message list on every `load`.
    ///
    /// The filter runs after raw messages are read from the store and before
    /// they are returned to the agent. Use it for truncation, summarization, or
    /// any other shaping. For reusable named policies, depend on `rig-memory`.
    pub fn with_filter<F>(mut self, filter: F) -> Self
    where
        F: MessageFilter + 'static,
    {
        self.filter = Some(Arc::new(filter));
        self
    }

    fn lock(
        &self,
    ) -> Result<std::sync::MutexGuard<'_, HashMap<String, Vec<Message>>>, MemoryError> {
        self.inner
            .lock()
            .map_err(|e| MemoryError::Internal(e.to_string()))
    }
}

impl std::fmt::Debug for InMemoryConversationMemory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("InMemoryConversationMemory")
            .field("filter", &self.filter.as_ref().map(|_| "<filter>"))
            .finish()
    }
}

impl ConversationMemory for InMemoryConversationMemory {
    fn load<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
        Box::pin(async move {
            let messages = {
                let guard = self.lock()?;
                guard.get(conversation_id).cloned().unwrap_or_default()
            };
            match &self.filter {
                Some(filter) => Ok(filter(messages)),
                None => Ok(messages),
            }
        })
    }

    fn append<'a>(
        &'a self,
        conversation_id: &'a str,
        messages: Vec<Message>,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        Box::pin(async move {
            let mut guard = self.lock()?;
            guard
                .entry(conversation_id.to_string())
                .or_default()
                .extend(messages);
            Ok(())
        })
    }

    fn clear<'a>(
        &'a self,
        conversation_id: &'a str,
    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
        Box::pin(async move {
            let mut guard = self.lock()?;
            guard.remove(conversation_id);
            Ok(())
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::completion::Message;

    fn user(text: &str) -> Message {
        Message::user(text)
    }

    fn assistant(text: &str) -> Message {
        Message::assistant(text)
    }

    #[tokio::test]
    async fn round_trip() {
        let mem = InMemoryConversationMemory::new();
        assert!(mem.load("c1").await.unwrap().is_empty());

        mem.append("c1", vec![user("hello"), assistant("hi")])
            .await
            .unwrap();

        let loaded = mem.load("c1").await.unwrap();
        assert_eq!(loaded.len(), 2);
    }

    #[tokio::test]
    async fn isolation_between_conversations() {
        let mem = InMemoryConversationMemory::new();
        mem.append("a", vec![user("hi a")]).await.unwrap();
        mem.append("b", vec![user("hi b")]).await.unwrap();

        assert_eq!(mem.load("a").await.unwrap().len(), 1);
        assert_eq!(mem.load("b").await.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn clear_removes_history() {
        let mem = InMemoryConversationMemory::new();
        mem.append("c", vec![user("x")]).await.unwrap();
        mem.clear("c").await.unwrap();
        assert!(mem.load("c").await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn with_filter_transforms_loaded_messages() {
        let mem = InMemoryConversationMemory::new()
            .with_filter(|msgs: Vec<Message>| msgs.into_iter().rev().take(2).collect());

        mem.append(
            "c",
            vec![user("1"), assistant("2"), user("3"), assistant("4")],
        )
        .await
        .unwrap();

        let loaded = mem.load("c").await.unwrap();
        assert_eq!(loaded.len(), 2, "filter should retain only 2 messages");
    }

    #[tokio::test]
    async fn arc_conversation_memory_forwards_to_inner() {
        let inner = Arc::new(InMemoryConversationMemory::new());
        let mem: Arc<dyn ConversationMemory> = inner.clone();

        mem.append("c", vec![user("hello")]).await.unwrap();

        assert_eq!(inner.load("c").await.unwrap().len(), 1);
        mem.clear("c").await.unwrap();
        assert!(inner.load("c").await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn boxed_conversation_memory_forwards_to_inner() {
        let mem: Box<dyn ConversationMemory> = Box::new(InMemoryConversationMemory::new());

        mem.append("c", vec![user("hello")]).await.unwrap();

        assert_eq!(mem.load("c").await.unwrap().len(), 1);
        mem.clear("c").await.unwrap();
        assert!(mem.load("c").await.unwrap().is_empty());
    }
}


/*
 * Copyright 2024 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 tool

import (
	"context"

	"github.com/cloudwego/eino/schema"
)

// BaseTool provides the metadata that a ChatModel uses to decide whether and
// how to call a tool. Info returns a [schema.ToolInfo] containing the tool
// name, description, and parameter JSON schema.
//
// BaseTool alone is sufficient when passing tool definitions to a ChatModel
// via WithTools — the model only needs the schema to generate tool calls.
// To also execute the tool, implement [InvokableTool] or [StreamableTool].
type BaseTool interface {
	Info(ctx context.Context) (*schema.ToolInfo, error)
}

// InvokableTool is a tool that can be executed by ToolsNode.
//
// InvokableRun receives the model's tool call arguments as a JSON-encoded
// string and returns a plain string result that is sent back to the model as
// a tool message. The framework handles JSON decoding automatically when using
// the [utils.InferTool] or [utils.NewTool] constructors.
type InvokableTool interface {
	BaseTool

	// InvokableRun executes the tool with arguments encoded as a JSON string.
	InvokableRun(ctx context.Context, argumentsInJSON string, opts ...Option) (string, error)
}

// StreamableTool is a streaming variant of [InvokableTool].
//
// StreamableRun returns a [schema.StreamReader] that yields string chunks
// incrementally. The caller (ToolsNode) is responsible for closing the reader.
type StreamableTool interface {
	BaseTool

	StreamableRun(ctx context.Context, argumentsInJSON string, opts ...Option) (*schema.StreamReader[string], error)
}

// EnhancedInvokableTool is a tool that returns structured multimodal results.
//
// Unlike [InvokableTool], arguments arrive as a [schema.ToolArgument] (not a
// raw JSON string) and the result is a [schema.ToolResult] which can carry
// text, images, audio, video, and file content.
//
// When a tool implements both a standard and an enhanced interface, ToolsNode
// prioritises the enhanced interface.
type EnhancedInvokableTool interface {
	BaseTool
	InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...Option) (*schema.ToolResult, error)
}

// EnhancedStreamableTool is the streaming variant of [EnhancedInvokableTool].
//
// It streams [schema.ToolResult] chunks, enabling incremental multimodal
// output. The caller is responsible for closing the returned [schema.StreamReader].
type EnhancedStreamableTool interface {
	BaseTool
	StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...Option) (*schema.StreamReader[*schema.ToolResult], error)
}


import { Command } from "commander";
import pc from "picocolors";
import ora from "ora";
import open from "open";
import boxen from "boxen";
import {
  saveTokens,
  clearTokens,
  getValidAccessToken,
  startDeviceAuthorization,
  pollDeviceToken,
  DEFAULT_DEVICE_POLL_INTERVAL_SECONDS,
} from "../utils/auth.js";

import { trackEvent } from "../utils/tracking.js";
import { CLI_CLIENT_ID } from "../constants.js";
import { getBaseUrl } from "../utils/api.js";

let baseUrl = "https://context7.com";

export function setAuthBaseUrl(url: string): void {
  baseUrl = url;
}

export function registerAuthCommands(program: Command): void {
  program
    .command("login")
    .description("Log in to Context7")
    .option("--no-browser", "Don't open browser automatically")
    .action(async (options) => {
      await loginCommand(options);
    });

  program
    .command("logout")
    .description("Log out of Context7")
    .action(() => {
      logoutCommand();
    });

  program
    .command("whoami")
    .description("Show current login status")
    .action(async () => {
      await whoamiCommand();
    });
}

function renderDeviceCodeBox(
  userCode: string,
  verificationUri: string,
  verificationUriComplete: string | undefined
): string {
  const codeLine = `${pc.dim("Your one-time code:")}\n\n    ${pc.green(pc.bold(userCode))}`;
  // Per RFC 8628 §3.3, even when verification_uri_complete is available we
  // still show the bare verification_uri so users on screen readers / paper
  // can type it manually.
  const linkLine = verificationUriComplete
    ? `${pc.dim("Open this link to approve:")}\n${pc.cyan(verificationUriComplete)}\n\n${pc.dim("Or visit")} ${pc.cyan(verificationUri)} ${pc.dim("and enter the code above.")}`
    : `${pc.dim("Visit:")} ${pc.cyan(verificationUri)}`;
  return boxen(`${codeLine}\n\n${linkLine}`, {
    title: "Sign in to Context7",
    titleAlignment: "left",
    padding: 1,
    margin: { top: 1, bottom: 1, left: 2, right: 2 },
    borderStyle: "round",
    borderColor: "gray",
  });
}

/** Prints a prompt and resolves on the next keypress. No-op when stdin isn't a TTY. */
function waitForEnter(prompt: string): Promise<void> {
  if (!process.stdin.isTTY) return Promise.resolve();
  return new Promise<void>((resolve) => {
    process.stdout.write(`  ${pc.dim(prompt)} `);
    const onData = (chunk: Buffer) => {
      // Ctrl-C
      if (chunk[0] === 0x03) {
        process.stdin.removeListener("data", onData);
        process.stdin.setRawMode?.(false);
        process.stdin.pause();
        process.stdout.write("\n");
        process.exit(130);
      }
      process.stdin.removeListener("data", onData);
      process.stdin.setRawMode?.(false);
      process.stdin.pause();
      process.stdout.write("\n");
      resolve();
    };
    process.stdin.setRawMode?.(true);
    process.stdin.resume();
    process.stdin.on("data", onData);
  });
}

async function announceIdentity(accessToken: string): Promise<string> {
  try {
    const whoami = await fetchWhoami(accessToken);
    const name = whoami.email || whoami.name;
    if (!name) return "Login successful!";
    const team = whoami.teamspace?.name;
    return team
      ? `Logged in as ${pc.bold(name)} ${pc.dim(`(${team})`)}`
      : `Logged in as ${pc.bold(name)}`;
  } catch {
    return "Login successful!";
  }
}

export async function performLogin(openBrowser = true): Promise<string | null> {
  const spinner = ora("Preparing login...").start();

  let authorization;
  try {
    authorization = await startDeviceAuthorization(baseUrl, CLI_CLIENT_ID);
  } catch (error) {
    spinner.fail(pc.red("Login failed"));
    if (error instanceof Error) console.error(pc.red(error.message));
    return null;
  }

  spinner.stop();

  console.log(
    renderDeviceCodeBox(
      authorization.user_code,
      authorization.verification_uri,
      authorization.verification_uri_complete
    )
  );

  const target = authorization.verification_uri_complete ?? authorization.verification_uri;
  if (openBrowser) {
    await waitForEnter("Press Enter to open the browser, or Ctrl-C to quit...");
    try {
      await open(target);
    } catch {
      console.log(pc.dim(`  Couldn't open a browser — visit the link above manually.`));
    }
  } else {
    console.log(pc.dim("  Open the link above in any browser to continue."));
    console.log("");
  }

  const waitingSpinner = ora({ text: "Waiting for authorization...", indent: 2 }).start();

  const deadline = Date.now() + authorization.expires_in * 1000;
  let intervalMs = (authorization.interval ?? DEFAULT_DEVICE_POLL_INTERVAL_SECONDS) * 1000;

  while (Date.now() < deadline) {
    await new Promise((resolve) => setTimeout(resolve, intervalMs));
    try {
      const result = await pollDeviceToken(baseUrl, CLI_CLIENT_ID, authorization.device_code);
      if (result.status === "approved" && result.tokens) {
        saveTokens(result.tokens);
        const successText = await announceIdentity(result.tokens.access_token);
        waitingSpinner.succeed(pc.green(successText));
        return result.tokens.access_token;
      }
      if (result.status === "slow_down") {
        intervalMs += 5000;
        continue;
      }
      if (result.status === "denied") {
        waitingSpinner.fail(pc.red("Authorization denied."));
        return null;
      }
      if (result.status === "expired") {
        waitingSpinner.fail(pc.red("Code expired. Run login again."));
        return null;
      }
      if (result.status === "transient") {
        // RFC 8628 §3.5: client MUST unilaterally reduce polling frequency on
        // connection timeout. Apply +5s like slow_down so a flaky network or
        // 5xx burst doesn't keep hitting at the original cadence.
        intervalMs += 5000;
        continue;
      }
      // pending — keep polling at the current cadence.
    } catch (error) {
      waitingSpinner.fail(pc.red("Login failed"));
      if (error instanceof Error) console.error(pc.red(error.message));
      return null;
    }
  }

  waitingSpinner.fail(pc.red("Code expired without approval."));
  return null;
}

async function loginCommand(options: { browser: boolean }): Promise<void> {
  trackEvent("command", { name: "login" });
  const existingToken = await getValidAccessToken();
  if (existingToken) {
    console.log(pc.yellow("You are already logged in."));
    console.log(pc.dim("Run 'ctx7 logout' first if you want to log in with a different account."));
    return;
  }
  clearTokens();

  const token = await performLogin(options.browser);
  if (!token) {
    process.exit(1);
  }
  console.log("");
  console.log(pc.dim("You can now use authenticated Context7 features."));
}

function logoutCommand(): void {
  trackEvent("command", { name: "logout" });
  if (clearTokens()) {
    console.log(pc.green("Logged out successfully."));
  } else {
    console.log(pc.yellow("You are not logged in."));
  }
}

async function whoamiCommand(): Promise<void> {
  trackEvent("command", { name: "whoami" });
  const accessToken = await getValidAccessToken();

  if (!accessToken) {
    console.log(pc.yellow("Not logged in."));
    console.log(pc.dim("Run 'ctx7 login' to authenticate."));
    return;
  }

  console.log(pc.green("Logged in"));

  try {
    const whoami = await fetchWhoami(accessToken);
    if (whoami.name) {
      console.log(`${pc.dim("Name:".padEnd(13))}${whoami.name}`);
    }
    if (whoami.email) {
      console.log(`${pc.dim("Email:".padEnd(13))}${whoami.email}`);
    }
    if (whoami.teamspace) {
      console.log(`${pc.dim("Teamspace:".padEnd(13))}${whoami.teamspace.name}`);
    }
  } catch {
    console.log(pc.dim("(Session may be expired - run 'ctx7 login' to refresh)"));
  }
}

interface WhoamiResponse {
  success: boolean;
  name: string | null;
  email: string | null;
  teamspace: { id: string; name: string } | null;
}

async function fetchWhoami(accessToken: string): Promise<WhoamiResponse> {
  const response = await fetch(`${getBaseUrl()}/api/dashboard/whoami`, {
    headers: {
      Authorization: `Bearer ${accessToken}`,
    },
  });

  if (!response.ok) {
    throw new Error("Failed to fetch user info");
  }

  return (await response.json()) as WhoamiResponse;
}