package session

import (
	"encoding/json"
	"sync"
	"time"
)

// State represents the current state of a session
type State int

const (
	Active State = iota
	Completed
	Killed
	TimedOut
)

func (s State) String() string {
	switch s {
	case Active:
		return "active"
	case Completed:
		return "completed"
	case Killed:
		return "killed"
	case TimedOut:
		return "timeout"
	default:
		return "unknown"
	}
}

// Session represents an agent session
type Session struct {
	mu sync.RWMutex

	ID           string            `json:"id"`
	State        State             `json:"state"`
	StartTime    time.Time         `json:"start_time"`
	LastActivity time.Time         `json:"last_activity"`
	EndTime      *time.Time        `json:"end_time,omitempty"`
	RequestCount int               `json:"request_count"`
	BytesIn      int64             `json:"bytes_in"`
	BytesOut     int64             `json:"bytes_out"`
	Backend      string            `json:"backend"`
	ClientAddr   string            `json:"client_addr"`
	Metadata     map[string]string `json:"metadata,omitempty"`

	// Track all backends used and request count per backend
	BackendsUsed map[string]int `json:"backends_used,omitempty"`

	// Terminated sessions cannot be resumed (for malicious agents)
	Terminated bool `json:"terminated,omitempty"`

	// Cached system prompt hash (avoids re-scanning identical prompts)
	SystemPromptHash string `json:"-"` // Not persisted, session-local cache

	// WebSocket session fields
	IsWebSocket  bool  `json:"is_websocket,omitempty"`
	FrameCount   int64 `json:"frame_count,omitempty"`
	TextFrames   int64 `json:"text_frames,omitempty"`
	BinaryFrames int64 `json:"binary_frames,omitempty"`

	// For rate limiting - track recent request times
	RequestTimes []time.Time `json:"-"`

	// For kill signaling
	killChan chan struct{}

	// Token tracking (for LLM API usage)
	TokensIn  int64 `json:"tokens_in"`
	TokensOut int64 `json:"tokens_out"`

	// Tool/function call tracking
	ToolCalls      int            `json:"tool_calls"`
	ToolCallCounts map[string]int `json:"tool_call_counts,omitempty"` // Per-tool counts

	// Tool call history (who called what)
	ToolCallHistory []ToolCallRecord `json:"tool_call_history,omitempty"`

	// Conversation history for failover/replay
	Messages     []Message `json:"messages,omitempty"`
	SystemPrompt string    `json:"system_prompt,omitempty"`

	// Failed backends for this session (for failover)
	FailedBackends []string `json:"failed_backends,omitempty"`
}

// ToolCallRecord tracks a single tool/function call
type ToolCallRecord struct {
	Timestamp time.Time      `json:"timestamp"`
	ToolName  string         `json:"tool_name"`
	ToolType  string         `json:"tool_type,omitempty"` // "function", "code_interpreter", etc.
	RequestID string         `json:"request_id,omitempty"`
	Arguments map[string]any `json:"arguments,omitempty"`
	Result    string         `json:"result,omitempty"`
}

// Message represents a single message in the conversation history
type Message struct {
	Role      string    `json:"role"` // "user", "assistant", "system"
	Content   string    `json:"content"`
	Timestamp time.Time `json:"timestamp"`
	Backend   string    `json:"backend,omitempty"` // which backend generated this
}

// SessionState is a serializable snapshot of session state for failover
type SessionState struct {
	ID             string            `json:"id"`
	StartTime      time.Time         `json:"start_time"`
	Messages       []Message         `json:"messages"`
	ToolCalls      []ToolCallRecord  `json:"tool_calls"`
	SystemPrompt   string            `json:"system_prompt,omitempty"`
	TokensIn       int64             `json:"tokens_in"`
	TokensOut      int64             `json:"tokens_out"`
	Metadata       map[string]string `json:"metadata,omitempty"`
	CurrentBackend string            `json:"current_backend"`
	FailedBackends []string          `json:"failed_backends,omitempty"`
}

// NewSession creates a new session with the given ID
func NewSession(id, backend, clientAddr string) *Session {
	now := time.Now()
	return &Session{
		ID:           id,
		State:        Active,
		StartTime:    now,
		LastActivity: now,
		Backend:      backend,
		ClientAddr:   clientAddr,
		Metadata:     make(map[string]string),
		BackendsUsed: make(map[string]int),
		killChan:     make(chan struct{}),
	}
}

// Touch updates the last activity time and increments request count
func (s *Session) Touch() {
	s.mu.Lock()
	defer s.mu.Unlock()
	now := time.Now()
	s.LastActivity = now
	s.RequestCount++

	// Track request time for rate limiting (keep last 2 minutes)
	s.RequestTimes = append(s.RequestTimes, now)
	cutoff := now.Add(-2 * time.Minute)
	for len(s.RequestTimes) > 0 && s.RequestTimes[0].Before(cutoff) {
		s.RequestTimes = s.RequestTimes[1:]
	}
}

// GetRequestTimes returns a copy of recent request times
func (s *Session) GetRequestTimes() []time.Time {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]time.Time, len(s.RequestTimes))
	copy(result, s.RequestTimes)
	return result
}

// AddBytes adds bytes to the session counters
func (s *Session) AddBytes(in, out int64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.BytesIn += in
	s.BytesOut += out
}

// TouchAndRecord atomically updates activity time, adds request bytes,
// and records the backend — all under a single lock acquisition.
// This prevents concurrent requests from observing inconsistent state
// between individual Touch/AddBytes/RecordBackend calls.
func (s *Session) TouchAndRecord(bytesIn int64, backend string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	s.LastActivity = now
	s.RequestCount++

	// Track request time for rate limiting (keep last 2 minutes)
	s.RequestTimes = append(s.RequestTimes, now)
	cutoff := now.Add(-2 * time.Minute)
	for len(s.RequestTimes) > 0 && s.RequestTimes[0].Before(cutoff) {
		s.RequestTimes = s.RequestTimes[1:]
	}

	s.BytesIn += bytesIn

	if s.BackendsUsed == nil {
		s.BackendsUsed = make(map[string]int)
	}
	s.BackendsUsed[backend]++
	s.Backend = backend
}

// AddTokens adds token counts to the session
func (s *Session) AddTokens(in, out int64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.TokensIn += in
	s.TokensOut += out
}

// GetTokens returns current token counts
func (s *Session) GetTokens() (in, out int64) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.TokensIn, s.TokensOut
}

// RecordToolCall records a tool/function call
func (s *Session) RecordToolCall(toolName, toolType, requestID, argsJSON string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.ToolCalls++

	if s.ToolCallCounts == nil {
		s.ToolCallCounts = make(map[string]int)
	}
	s.ToolCallCounts[toolName]++

	// Parse arguments JSON into map for structured access
	var args map[string]any
	if argsJSON != "" {
		_ = json.Unmarshal([]byte(argsJSON), &args)
	}

	// Keep history (limited to avoid memory bloat)
	record := ToolCallRecord{
		Timestamp: time.Now(),
		ToolName:  toolName,
		ToolType:  toolType,
		RequestID: requestID,
		Arguments: args,
	}
	s.ToolCallHistory = append(s.ToolCallHistory, record)
	if len(s.ToolCallHistory) > 100 {
		s.ToolCallHistory = s.ToolCallHistory[1:]
	}
}

// GetToolCalls returns total tool call count
func (s *Session) GetToolCalls() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.ToolCalls
}

// GetToolFanout returns the number of distinct tools used
func (s *Session) GetToolFanout() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.ToolCallCounts)
}

// GetToolCallCounts returns a copy of per-tool call counts
func (s *Session) GetToolCallCounts() map[string]int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make(map[string]int, len(s.ToolCallCounts))
	for k, v := range s.ToolCallCounts {
		result[k] = v
	}
	return result
}

// GetToolCallHistory returns a copy of the tool call history
func (s *Session) GetToolCallHistory() []ToolCallRecord {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]ToolCallRecord, len(s.ToolCallHistory))
	copy(result, s.ToolCallHistory)
	return result
}

// GetSystemPromptHash returns the cached system prompt hash
func (s *Session) GetSystemPromptHash() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.SystemPromptHash
}

// SetSystemPromptHash sets the cached system prompt hash
func (s *Session) SetSystemPromptHash(hash string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.SystemPromptHash = hash
}

// FrameType represents the type of WebSocket frame
type FrameType int

const (
	FrameText FrameType = iota
	FrameBinary
)

// FrameDirection represents the direction of a WebSocket frame
type FrameDirection int

const (
	FrameInbound FrameDirection = iota
	FrameOutbound
)

// AddFrame records a WebSocket frame and updates byte counters
func (s *Session) AddFrame(frameType FrameType, size int64, direction FrameDirection) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.FrameCount++
	s.LastActivity = time.Now()

	switch frameType {
	case FrameText:
		s.TextFrames++
	case FrameBinary:
		s.BinaryFrames++
	}

	switch direction {
	case FrameInbound:
		s.BytesIn += size
	case FrameOutbound:
		s.BytesOut += size
	}
}

// SetWebSocket marks this session as a WebSocket session
func (s *Session) SetWebSocket() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.IsWebSocket = true
}

// RecordBackend tracks which backend was used for a request
func (s *Session) RecordBackend(backend string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.BackendsUsed == nil {
		s.BackendsUsed = make(map[string]int)
	}
	s.BackendsUsed[backend]++
	// Also update the Backend field to show the last used backend
	s.Backend = backend
}

// GetBackendsUsed returns a copy of the backends used map
func (s *Session) GetBackendsUsed() map[string]int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make(map[string]int, len(s.BackendsUsed))
	for k, v := range s.BackendsUsed {
		result[k] = v
	}
	return result
}

// SetState updates the session state
func (s *Session) SetState(state State) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.State = state
	if state != Active {
		now := time.Now()
		s.EndTime = &now
	}
}

// GetState returns the current session state
func (s *Session) GetState() State {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.State
}

// IsActive returns true if the session is still active
func (s *Session) IsActive() bool {
	return s.GetState() == Active
}

// Kill signals the session to terminate (can be resumed later)
func (s *Session) Kill() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.State == Active {
		s.State = Killed
		now := time.Now()
		s.EndTime = &now
		close(s.killChan)
	}
}

// Terminate permanently kills the session (cannot be resumed)
// Use this for malicious or runaway agents
func (s *Session) Terminate() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.State == Active || s.State == Killed {
		s.State = Killed
		s.Terminated = true
		now := time.Now()
		s.EndTime = &now
		// Only close if not already closed
		select {
		case <-s.killChan:
			// Already closed
		default:
			close(s.killChan)
		}
	}
}

// IsTerminated returns true if the session was permanently terminated
func (s *Session) IsTerminated() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.Terminated
}

// Resume reactivates a killed session, allowing new requests
// Returns false if session is terminated (cannot be resumed)
func (s *Session) Resume() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.Terminated {
		return false // Cannot resume terminated sessions
	}
	if s.State == Killed {
		s.State = Active
		s.EndTime = nil
		s.LastActivity = time.Now()
		// Create new kill channel for future kill operations
		s.killChan = make(chan struct{})
		return true
	}
	return false
}

// KillChan returns the channel that's closed when the session is killed
func (s *Session) KillChan() <-chan struct{} {
	return s.killChan
}

// Duration returns how long the session has been running
func (s *Session) Duration() time.Duration {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.EndTime != nil {
		return s.EndTime.Sub(s.StartTime)
	}
	return time.Since(s.StartTime)
}

// IdleTime returns how long since the last activity
func (s *Session) IdleTime() time.Duration {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return time.Since(s.LastActivity)
}

// SetMetadata sets a metadata key-value pair
func (s *Session) SetMetadata(key, value string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Metadata[key] = value
}

// RecordMessage adds a message to the conversation history
func (s *Session) RecordMessage(role, content, backend string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	msg := Message{
		Role:      role,
		Content:   content,
		Timestamp: time.Now(),
		Backend:   backend,
	}

	s.Messages = append(s.Messages, msg)

	// Keep last 100 messages to avoid memory bloat
	if len(s.Messages) > 100 {
		s.Messages = s.Messages[1:]
	}
}

// SetSystemPrompt sets the system prompt for the session
func (s *Session) SetSystemPrompt(prompt string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.SystemPrompt = prompt
}

// GetMessages returns a copy of the conversation history
func (s *Session) GetMessages() []Message {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]Message, len(s.Messages))
	copy(result, s.Messages)
	return result
}

// AddFailedBackend records a backend that failed for this session
func (s *Session) AddFailedBackend(backend string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Check if already recorded
	for _, b := range s.FailedBackends {
		if b == backend {
			return
		}
	}
	s.FailedBackends = append(s.FailedBackends, backend)
}

// GetFailedBackends returns the list of failed backends
func (s *Session) GetFailedBackends() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]string, len(s.FailedBackends))
	copy(result, s.FailedBackends)
	return result
}

// Serialize creates a serializable snapshot of session state for failover
func (s *Session) Serialize() *SessionState {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// Copy messages
	messages := make([]Message, len(s.Messages))
	copy(messages, s.Messages)

	// Copy tool calls
	toolCalls := make([]ToolCallRecord, len(s.ToolCallHistory))
	copy(toolCalls, s.ToolCallHistory)

	// Copy metadata
	metadata := make(map[string]string, len(s.Metadata))
	for k, v := range s.Metadata {
		metadata[k] = v
	}

	// Copy failed backends
	failedBackends := make([]string, len(s.FailedBackends))
	copy(failedBackends, s.FailedBackends)

	return &SessionState{
		ID:             s.ID,
		StartTime:      s.StartTime,
		Messages:       messages,
		ToolCalls:      toolCalls,
		SystemPrompt:   s.SystemPrompt,
		TokensIn:       s.TokensIn,
		TokensOut:      s.TokensOut,
		Metadata:       metadata,
		CurrentBackend: s.Backend,
		FailedBackends: failedBackends,
	}
}

// Snapshot returns a copy of the session for safe reading.
// Note: Returns a new Session with a fresh zero-value mutex, not copying s.mu.
//
//nolint:govet // mutex field is zero-initialized, not copied from source
func (s *Session) Snapshot() Session {
	s.mu.RLock()
	defer s.mu.RUnlock()

	snap := Session{
		ID:           s.ID,
		State:        s.State,
		StartTime:    s.StartTime,
		LastActivity: s.LastActivity,
		EndTime:      s.EndTime,
		RequestCount: s.RequestCount,
		BytesIn:      s.BytesIn,
		BytesOut:     s.BytesOut,
		Backend:      s.Backend,
		ClientAddr:   s.ClientAddr,
		Metadata:     make(map[string]string, len(s.Metadata)),
		BackendsUsed: make(map[string]int, len(s.BackendsUsed)),
		IsWebSocket:  s.IsWebSocket,
		FrameCount:   s.FrameCount,
		TextFrames:   s.TextFrames,
		BinaryFrames: s.BinaryFrames,
		TokensIn:     s.TokensIn,
		TokensOut:    s.TokensOut,
		ToolCalls:    s.ToolCalls,
	}
	// Snapshots are read-only; a closed killChan prevents nil-channel blocks
	snap.killChan = make(chan struct{})
	close(snap.killChan)
	for k, v := range s.Metadata {
		snap.Metadata[k] = v
	}
	for k, v := range s.BackendsUsed {
		snap.BackendsUsed[k] = v
	}
	if s.ToolCallCounts != nil {
		snap.ToolCallCounts = make(map[string]int, len(s.ToolCallCounts))
		for k, v := range s.ToolCallCounts {
			snap.ToolCallCounts[k] = v
		}
	}
	if s.ToolCallHistory != nil {
		snap.ToolCallHistory = make([]ToolCallRecord, len(s.ToolCallHistory))
		copy(snap.ToolCallHistory, s.ToolCallHistory)
	}
	if s.Messages != nil {
		snap.Messages = make([]Message, len(s.Messages))
		copy(snap.Messages, s.Messages)
	}
	snap.SystemPrompt = s.SystemPrompt
	if s.FailedBackends != nil {
		snap.FailedBackends = make([]string, len(s.FailedBackends))
		copy(snap.FailedBackends, s.FailedBackends)
	}
	return snap
}
