package ports

import (
	"database/sql"
	"time"

	_ "github.com/jackc/pgx/v4/stdlib"
)

const DBSchema = `
CREATE TABLE IF NOT EXISTS port_attempts(
	id serial PRIMARY KEY,
	date timestamptz,
	remote_ip inet,
	port varchar(5)
);`

type PostgresStore struct {
	db *sql.DB
}

func NewPostgresStore(dsn string) (*PostgresStore, error) {
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		return nil, err
	}

	s := &PostgresStore{db: db}
	return s, nil
}

func (s *PostgresStore) InitDB() error {
	_, err := s.db.Exec(DBSchema)
	return err
}

func (s *PostgresStore) Add(attempt *ConnectionAttempt) error {
	stmt := `INSERT INTO port_attempts(date, remote_ip, port)
	VALUES ($1, $2, $3)`
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	_, err = tx.Exec(stmt, time.Now(), attempt.From, attempt.Port)
	if err != nil {
		return err
	}
	return tx.Commit()
}

func (s *PostgresStore) List() ([]*ConnectionAttempt, error) {
	stmt := `SELECT remote_ip, port FROM port_attempts`
	rows, err := s.db.Query(stmt)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var attempts []*ConnectionAttempt
	for rows.Next() {
		var a ConnectionAttempt
		if err := rows.Scan(&a.From, &a.Port); err != nil {
			return nil, err
		}
		attempts = append(attempts, &a)
	}

	return attempts, nil
}
