path: root/internal/store/agents.go

package store

import (
	"database/sql"
	"errors"
	"time"

	"task-dashboard/internal/models"
)

// CreateAgentSession creates a new pending agent session
func (s *Store) CreateAgentSession(session *models.AgentSession) error {
	result, err := s.db.Exec(`
		INSERT INTO agent_sessions (request_token, agent_name, agent_id, status, expires_at)
		VALUES (?, ?, ?, 'pending', ?)
	`, session.RequestToken, session.AgentName, session.AgentID, session.ExpiresAt)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}
	session.ID = id
	return nil
}

// GetAgentSessionByRequestToken retrieves a session by request token
func (s *Store) GetAgentSessionByRequestToken(token string) (*models.AgentSession, error) {
	var session models.AgentSession
	var sessionToken sql.NullString
	var sessionExpiresAt sql.NullTime

	err := s.db.QueryRow(`
		SELECT id, request_token, agent_name, agent_id, status, created_at, expires_at, session_token, session_expires_at
		FROM agent_sessions
		WHERE request_token = ?
	`, token).Scan(
		&session.ID,
		&session.RequestToken,
		&session.AgentName,
		&session.AgentID,
		&session.Status,
		&session.CreatedAt,
		&session.ExpiresAt,
		&sessionToken,
		&sessionExpiresAt,
	)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	if sessionToken.Valid {
		session.SessionToken = sessionToken.String
	}
	if sessionExpiresAt.Valid {
		session.SessionExpiresAt = &sessionExpiresAt.Time
	}
	return &session, nil
}

// GetPendingAgentSessionByAgentID retrieves an existing pending session for an agent
func (s *Store) GetPendingAgentSessionByAgentID(agentID string) (*models.AgentSession, error) {
	var session models.AgentSession

	err := s.db.QueryRow(`
		SELECT id, request_token, agent_name, agent_id, status, created_at, expires_at
		FROM agent_sessions
		WHERE agent_id = ? AND status = 'pending' AND expires_at > datetime('now', 'localtime')
		ORDER BY created_at DESC
		LIMIT 1
	`, agentID).Scan(
		&session.ID,
		&session.RequestToken,
		&session.AgentName,
		&session.AgentID,
		&session.Status,
		&session.CreatedAt,
		&session.ExpiresAt,
	)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	return &session, nil
}

// GetAgentSessionBySessionToken retrieves a session by session token
func (s *Store) GetAgentSessionBySessionToken(token string) (*models.AgentSession, error) {
	var session models.AgentSession
	var sessionToken sql.NullString
	var sessionExpiresAt sql.NullTime

	err := s.db.QueryRow(`
		SELECT id, request_token, agent_name, agent_id, status, created_at, expires_at, session_token, session_expires_at
		FROM agent_sessions
		WHERE session_token = ? AND status = 'approved'
	`, token).Scan(
		&session.ID,
		&session.RequestToken,
		&session.AgentName,
		&session.AgentID,
		&session.Status,
		&session.CreatedAt,
		&session.ExpiresAt,
		&sessionToken,
		&sessionExpiresAt,
	)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	if sessionToken.Valid {
		session.SessionToken = sessionToken.String
	}
	if sessionExpiresAt.Valid {
		session.SessionExpiresAt = &sessionExpiresAt.Time
	}
	return &session, nil
}

// ApproveAgentSession approves a pending session
func (s *Store) ApproveAgentSession(requestToken, sessionToken string, sessionExpiresAt time.Time) error {
	result, err := s.db.Exec(`
		UPDATE agent_sessions
		SET status = 'approved', session_token = ?, session_expires_at = ?
		WHERE request_token = ? AND status = 'pending'
	`, sessionToken, sessionExpiresAt, requestToken)
	if err != nil {
		return err
	}
	affected, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if affected == 0 {
		return errors.New("session not found or already processed")
	}
	return nil
}

// DenyAgentSession denies a pending session
func (s *Store) DenyAgentSession(requestToken string) error {
	result, err := s.db.Exec(`
		UPDATE agent_sessions
		SET status = 'denied'
		WHERE request_token = ? AND status = 'pending'
	`, requestToken)
	if err != nil {
		return err
	}
	affected, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if affected == 0 {
		return errors.New("session not found or already processed")
	}
	return nil
}

// GetPendingAgentSessions retrieves all unexpired pending sessions
func (s *Store) GetPendingAgentSessions() ([]models.AgentSession, error) {
	rows, err := s.db.Query(`
		SELECT id, request_token, agent_name, agent_id, status, created_at, expires_at
		FROM agent_sessions
		WHERE status = 'pending' AND expires_at > datetime('now', 'localtime')
		ORDER BY created_at DESC
	`)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()

	var sessions []models.AgentSession
	for rows.Next() {
		var session models.AgentSession
		if err := rows.Scan(
			&session.ID,
			&session.RequestToken,
			&session.AgentName,
			&session.AgentID,
			&session.Status,
			&session.CreatedAt,
			&session.ExpiresAt,
		); err != nil {
			return nil, err
		}
		sessions = append(sessions, session)
	}
	return sessions, rows.Err()
}

// InvalidatePreviousAgentSessions marks previous sessions for an agent as expired
func (s *Store) InvalidatePreviousAgentSessions(agentID string) error {
	_, err := s.db.Exec(`
		UPDATE agent_sessions
		SET status = 'expired'
		WHERE agent_id = ? AND status IN ('pending', 'approved')
	`, agentID)
	return err
}

// GetAgentByAgentID retrieves an agent by their agent_id (UUID)
func (s *Store) GetAgentByAgentID(agentID string) (*models.Agent, error) {
	var agent models.Agent
	var lastSeen sql.NullTime

	err := s.db.QueryRow(`
		SELECT id, name, agent_id, created_at, last_seen, trusted
		FROM agents
		WHERE agent_id = ?
	`, agentID).Scan(
		&agent.ID,
		&agent.Name,
		&agent.AgentID,
		&agent.CreatedAt,
		&lastSeen,
		&agent.Trusted,
	)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	if lastSeen.Valid {
		agent.LastSeen = &lastSeen.Time
	}
	return &agent, nil
}

// GetAgentByName retrieves an agent by name
func (s *Store) GetAgentByName(name string) (*models.Agent, error) {
	var agent models.Agent
	var lastSeen sql.NullTime

	err := s.db.QueryRow(`
		SELECT id, name, agent_id, created_at, last_seen, trusted
		FROM agents
		WHERE name = ?
	`, name).Scan(
		&agent.ID,
		&agent.Name,
		&agent.AgentID,
		&agent.CreatedAt,
		&lastSeen,
		&agent.Trusted,
	)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	if lastSeen.Valid {
		agent.LastSeen = &lastSeen.Time
	}
	return &agent, nil
}

// CreateOrUpdateAgent creates or updates an agent record
func (s *Store) CreateOrUpdateAgent(name, agentID string) error {
	_, err := s.db.Exec(`
		INSERT INTO agents (name, agent_id, last_seen, trusted)
		VALUES (?, ?, datetime('now'), 1)
		ON CONFLICT(agent_id) DO UPDATE SET
			name = excluded.name,
			last_seen = datetime('now')
	`, name, agentID)
	return err
}

// UpdateAgentLastSeen updates the last_seen timestamp for an agent
func (s *Store) UpdateAgentLastSeen(agentID string) error {
	_, err := s.db.Exec(`
		UPDATE agents SET last_seen = datetime('now')
		WHERE agent_id = ?
	`, agentID)
	return err
}

// GetAllAgents retrieves all agents
func (s *Store) GetAllAgents() ([]models.Agent, error) {
	rows, err := s.db.Query(`
		SELECT id, name, agent_id, created_at, last_seen, trusted
		FROM agents
		ORDER BY last_seen DESC NULLS LAST
	`)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()

	var agents []models.Agent
	for rows.Next() {
		var agent models.Agent
		var lastSeen sql.NullTime
		if err := rows.Scan(
			&agent.ID,
			&agent.Name,
			&agent.AgentID,
			&agent.CreatedAt,
			&lastSeen,
			&agent.Trusted,
		); err != nil {
			return nil, err
		}
		if lastSeen.Valid {
			agent.LastSeen = &lastSeen.Time
		}
		agents = append(agents, agent)
	}
	return agents, rows.Err()
}

// RevokeAgent sets trusted=false for an agent
func (s *Store) RevokeAgent(agentID string) error {
	_, err := s.db.Exec(`UPDATE agents SET trusted = 0 WHERE agent_id = ?`, agentID)
	return err
}

// CheckAgentTrust determines trust level for an agent request
func (s *Store) CheckAgentTrust(name, agentID string) (models.AgentTrustLevel, error) {
	// Check if this exact agent_id is known
	existingByID, err := s.GetAgentByAgentID(agentID)
	if err != nil {
		return "", err
	}

	// Check if this name is known with a different ID
	existingByName, err := s.GetAgentByName(name)
	if err != nil {
		return "", err
	}

	if existingByID != nil && existingByID.Name == name && existingByID.Trusted {
		return models.AgentTrustRecognized, nil
	}

	if existingByName != nil && existingByName.AgentID != agentID {
		return models.AgentTrustSuspicious, nil
	}

	return models.AgentTrustNew, nil
}