package database

import (
	"ai-code-assistant/pkg/config"
	"context"
	"embed"
	"errors"
	"fmt"
	"github.com/golang-migrate/migrate/v4"
	_ "github.com/golang-migrate/migrate/v4/database/postgres"
	"github.com/golang-migrate/migrate/v4/source/iofs"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	_ "github.com/lib/pq"
	"strconv"
)

//go:embed migrations
var migrations embed.FS

type contextKey string

const (
	contextKeyDB contextKey = "db"
)

func preparedStatements() map[string]string {
	return map[string]string{
		"insert_repo":    `INSERT INTO repos (repo_id, repo_hash, repo_path) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING`,
		"repo_from_path": `SELECT repo_id FROM repos WHERE repo_path = $1`,
		"get_chunk":      `SELECT document FROM langchain_pg_embedding WHERE JSON_EXTRACT_PATH_TEXT(cmetadata, 'chunk_id')=$1 AND JSON_EXTRACT_PATH_TEXT(cmetadata, 'file_path')=$2 AND JSON_EXTRACT_PATH_TEXT(cmetadata, 'repo_id')=$3`,
	}
}

type Database struct {
	pool *pgxpool.Pool
}

func (db *Database) DB(ctx context.Context) (*pgxpool.Conn, error) {
	conn, err := db.pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}

	for nm, query := range preparedStatements() {
		if _, err := conn.Conn().Prepare(ctx, nm, query); err != nil {
			return nil, fmt.Errorf("problem preparing statement %s: %w", nm, err)
		}
	}

	return conn, nil
}

func (db *Database) RepoIDFromPath(ctx context.Context, path string) (string, error) {
	conn, err := db.DB(ctx)
	if err != nil {
		return "", err
	}
	defer conn.Release()

	var repoID string

	if err := conn.QueryRow(ctx, "repo_from_path", path).Scan(&repoID); err != nil {
		return "", err
	}

	return repoID, nil
}

func (db *Database) GetChunk(ctx context.Context, chunkID int, path, repoID string) (string, error) {
	conn, err := db.DB(ctx)
	if err != nil {
		return "", err
	}
	defer conn.Release()

	var chunk string

	chunkIDStr := strconv.FormatInt(int64(chunkID), 10)

	if err := conn.QueryRow(ctx, "get_chunk", chunkIDStr, path, repoID).Scan(&chunk); err != nil {
		return "", err
	}

	return chunk, nil
}

func (db *Database) GetChunkContext(ctx context.Context, chunkID, distance int, path, repoID string) (string, error) {
	minChunk := chunkID - distance
	if minChunk < 0 {
		minChunk = 0
	}

	chunkContext := ""

	for chunkID := minChunk; chunkID < chunkID+(distance*2); chunkID++ {
		chunk, err := db.GetChunk(ctx, chunkID, path, repoID)
		if err == nil {
			chunkContext += chunk
		} else if !errors.Is(err, pgx.ErrNoRows) {
			return "", err
		} else {
			break
		}
	}

	return chunkContext, nil
}

func FromConfig(ctx context.Context, cfg *config.Configuration) (*Database, error) {
	migFS, err := iofs.New(migrations, "migrations")
	if err != nil {
		return nil, err
	}

	mig, err := migrate.NewWithSourceInstance("iofs", migFS, cfg.Database.ConnString)
	if err != nil {
		return nil, err
	}

	if err := mig.Up(); err != nil && !errors.Is(err, migrate.ErrNoChange) {
		return nil, fmt.Errorf("problem performing database migrations: %w", err)
	}

	pool, err := pgxpool.New(ctx, cfg.Database.ConnString)
	if err != nil {
		return nil, err
	}

	return &Database{pool: pool}, nil
}

func FromContext(ctx context.Context) *Database {
	return ctx.Value(contextKeyDB).(*Database)
}

func WrapContext(ctx context.Context, cfg *Database) context.Context {
	return context.WithValue(ctx, contextKeyDB, cfg)
}