package investigation

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
)

// ClickHouseStore keeps the latest finished record per (ip, mode) as JSON.
// No TTL: what an address turned out to be is worth keeping.
type ClickHouseStore struct {
	conn driver.Conn
	db   string
}

func NewClickHouseStore(ctx context.Context, conn driver.Conn, db string) (*ClickHouseStore, error) {
	err := conn.Exec(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.ip_investigations (
		ip String,
		mode LowCardinality(String),
		status LowCardinality(String),
		started_at DateTime64(3, 'UTC'),
		finished_at DateTime64(3, 'UTC'),
		record String
	) ENGINE = ReplacingMergeTree(finished_at)
	ORDER BY (ip, mode)`, db))
	if err != nil {
		return nil, err
	}
	return &ClickHouseStore{conn: conn, db: db}, nil
}

func (s *ClickHouseStore) Save(ctx context.Context, rec Record) error {
	body, err := json.Marshal(rec)
	if err != nil {
		return err
	}
	finished := time.Now().UTC()
	if rec.FinishedAt != nil {
		finished = *rec.FinishedAt
	}
	return s.conn.Exec(ctx, fmt.Sprintf(`INSERT INTO %s.ip_investigations (ip, mode, status, started_at, finished_at, record) VALUES (?, ?, ?, ?, ?, ?)`, s.db),
		rec.IP, rec.Mode, rec.Status, rec.StartedAt.UTC(), finished, string(body))
}

func (s *ClickHouseStore) Load(ctx context.Context, ip string) ([]Record, error) {
	rows, err := s.conn.Query(ctx, fmt.Sprintf(`SELECT record FROM %s.ip_investigations FINAL WHERE ip = ?`, s.db), ip)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Record
	for rows.Next() {
		var body string
		if err := rows.Scan(&body); err != nil {
			return nil, err
		}
		var rec Record
		if err := json.Unmarshal([]byte(body), &rec); err == nil {
			out = append(out, rec)
		}
	}
	return out, rows.Err()
}