package db

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

	"github.com/gooseek/backend/internal/computer"
)

type ComputerTaskRepo struct {
	db *sql.DB
}

func NewComputerTaskRepo(db *sql.DB) *ComputerTaskRepo {
	return &ComputerTaskRepo{db: db}
}

func (r *ComputerTaskRepo) Migrate() error {
	query := `
	CREATE TABLE IF NOT EXISTS computer_tasks (
		id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
		user_id UUID NOT NULL,
		query TEXT NOT NULL,
		status VARCHAR(20) NOT NULL DEFAULT 'pending',
		plan JSONB,
		sub_tasks JSONB,
		artifacts JSONB,
		memory JSONB,
		progress INT DEFAULT 0,
		message TEXT,
		error TEXT,
		schedule JSONB,
		next_run_at TIMESTAMPTZ,
		run_count INT DEFAULT 0,
		total_cost DECIMAL(10,6) DEFAULT 0,
		created_at TIMESTAMPTZ DEFAULT NOW(),
		updated_at TIMESTAMPTZ DEFAULT NOW(),
		completed_at TIMESTAMPTZ
	);

	CREATE INDEX IF NOT EXISTS idx_computer_tasks_user_id ON computer_tasks(user_id);
	CREATE INDEX IF NOT EXISTS idx_computer_tasks_status ON computer_tasks(status);
	CREATE INDEX IF NOT EXISTS idx_computer_tasks_next_run ON computer_tasks(next_run_at) WHERE next_run_at IS NOT NULL;
	CREATE INDEX IF NOT EXISTS idx_computer_tasks_created ON computer_tasks(created_at DESC);
	`

	_, err := r.db.Exec(query)
	return err
}

func (r *ComputerTaskRepo) Create(ctx context.Context, task *computer.ComputerTask) error {
	planJSON, _ := json.Marshal(task.Plan)
	subTasksJSON, _ := json.Marshal(task.SubTasks)
	artifactsJSON, _ := json.Marshal(task.Artifacts)
	memoryJSON, _ := json.Marshal(task.Memory)
	scheduleJSON, _ := json.Marshal(task.Schedule)

	query := `
		INSERT INTO computer_tasks (
			id, user_id, query, status, plan, sub_tasks, artifacts, memory,
			progress, message, error, schedule, next_run_at, run_count, total_cost,
			created_at, updated_at, completed_at
		) VALUES (
			$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18
		)
	`

	_, err := r.db.ExecContext(ctx, query,
		task.ID,
		task.UserID,
		task.Query,
		task.Status,
		planJSON,
		subTasksJSON,
		artifactsJSON,
		memoryJSON,
		task.Progress,
		task.Message,
		task.Error,
		scheduleJSON,
		task.NextRunAt,
		task.RunCount,
		task.TotalCost,
		task.CreatedAt,
		task.UpdatedAt,
		task.CompletedAt,
	)

	return err
}

func (r *ComputerTaskRepo) Update(ctx context.Context, task *computer.ComputerTask) error {
	planJSON, _ := json.Marshal(task.Plan)
	subTasksJSON, _ := json.Marshal(task.SubTasks)
	artifactsJSON, _ := json.Marshal(task.Artifacts)
	memoryJSON, _ := json.Marshal(task.Memory)
	scheduleJSON, _ := json.Marshal(task.Schedule)

	query := `
		UPDATE computer_tasks SET
			status = $1,
			plan = $2,
			sub_tasks = $3,
			artifacts = $4,
			memory = $5,
			progress = $6,
			message = $7,
			error = $8,
			schedule = $9,
			next_run_at = $10,
			run_count = $11,
			total_cost = $12,
			updated_at = $13,
			completed_at = $14
		WHERE id = $15
	`

	_, err := r.db.ExecContext(ctx, query,
		task.Status,
		planJSON,
		subTasksJSON,
		artifactsJSON,
		memoryJSON,
		task.Progress,
		task.Message,
		task.Error,
		scheduleJSON,
		task.NextRunAt,
		task.RunCount,
		task.TotalCost,
		time.Now(),
		task.CompletedAt,
		task.ID,
	)

	return err
}

func (r *ComputerTaskRepo) GetByID(ctx context.Context, id string) (*computer.ComputerTask, error) {
	query := `
		SELECT id, user_id, query, status, plan, sub_tasks, artifacts, memory,
			progress, message, error, schedule, next_run_at, run_count, total_cost,
			created_at, updated_at, completed_at
		FROM computer_tasks
		WHERE id = $1
	`

	var task computer.ComputerTask
	var planJSON, subTasksJSON, artifactsJSON, memoryJSON, scheduleJSON []byte
	var message, errStr sql.NullString
	var nextRunAt, completedAt sql.NullTime

	err := r.db.QueryRowContext(ctx, query, id).Scan(
		&task.ID,
		&task.UserID,
		&task.Query,
		&task.Status,
		&planJSON,
		&subTasksJSON,
		&artifactsJSON,
		&memoryJSON,
		&task.Progress,
		&message,
		&errStr,
		&scheduleJSON,
		&nextRunAt,
		&task.RunCount,
		&task.TotalCost,
		&task.CreatedAt,
		&task.UpdatedAt,
		&completedAt,
	)

	if err != nil {
		return nil, err
	}

	if len(planJSON) > 0 {
		json.Unmarshal(planJSON, &task.Plan)
	}
	if len(subTasksJSON) > 0 {
		json.Unmarshal(subTasksJSON, &task.SubTasks)
	}
	if len(artifactsJSON) > 0 {
		json.Unmarshal(artifactsJSON, &task.Artifacts)
	}
	if len(memoryJSON) > 0 {
		json.Unmarshal(memoryJSON, &task.Memory)
	}
	if len(scheduleJSON) > 0 {
		json.Unmarshal(scheduleJSON, &task.Schedule)
	}

	if message.Valid {
		task.Message = message.String
	}
	if errStr.Valid {
		task.Error = errStr.String
	}
	if nextRunAt.Valid {
		task.NextRunAt = &nextRunAt.Time
	}
	if completedAt.Valid {
		task.CompletedAt = &completedAt.Time
	}

	return &task, nil
}

func (r *ComputerTaskRepo) GetByUserID(ctx context.Context, userID string, limit, offset int) ([]computer.ComputerTask, error) {
	query := `
		SELECT id, user_id, query, status, plan, sub_tasks, artifacts, memory,
			progress, message, error, schedule, next_run_at, run_count, total_cost,
			created_at, updated_at, completed_at
		FROM computer_tasks
		WHERE user_id = $1
		ORDER BY created_at DESC
		LIMIT $2 OFFSET $3
	`

	rows, err := r.db.QueryContext(ctx, query, userID, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []computer.ComputerTask

	for rows.Next() {
		var task computer.ComputerTask
		var planJSON, subTasksJSON, artifactsJSON, memoryJSON, scheduleJSON []byte
		var message, errStr sql.NullString
		var nextRunAt, completedAt sql.NullTime

		err := rows.Scan(
			&task.ID,
			&task.UserID,
			&task.Query,
			&task.Status,
			&planJSON,
			&subTasksJSON,
			&artifactsJSON,
			&memoryJSON,
			&task.Progress,
			&message,
			&errStr,
			&scheduleJSON,
			&nextRunAt,
			&task.RunCount,
			&task.TotalCost,
			&task.CreatedAt,
			&task.UpdatedAt,
			&completedAt,
		)
		if err != nil {
			continue
		}

		if len(planJSON) > 0 {
			json.Unmarshal(planJSON, &task.Plan)
		}
		if len(subTasksJSON) > 0 {
			json.Unmarshal(subTasksJSON, &task.SubTasks)
		}
		if len(artifactsJSON) > 0 {
			json.Unmarshal(artifactsJSON, &task.Artifacts)
		}
		if len(memoryJSON) > 0 {
			json.Unmarshal(memoryJSON, &task.Memory)
		}
		if len(scheduleJSON) > 0 {
			json.Unmarshal(scheduleJSON, &task.Schedule)
		}

		if message.Valid {
			task.Message = message.String
		}
		if errStr.Valid {
			task.Error = errStr.String
		}
		if nextRunAt.Valid {
			task.NextRunAt = &nextRunAt.Time
		}
		if completedAt.Valid {
			task.CompletedAt = &completedAt.Time
		}

		tasks = append(tasks, task)
	}

	return tasks, nil
}

func (r *ComputerTaskRepo) GetScheduled(ctx context.Context) ([]computer.ComputerTask, error) {
	query := `
		SELECT id, user_id, query, status, plan, sub_tasks, artifacts, memory,
			progress, message, error, schedule, next_run_at, run_count, total_cost,
			created_at, updated_at, completed_at
		FROM computer_tasks
		WHERE status = 'scheduled' AND schedule IS NOT NULL
		ORDER BY next_run_at ASC
	`

	rows, err := r.db.QueryContext(ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []computer.ComputerTask

	for rows.Next() {
		var task computer.ComputerTask
		var planJSON, subTasksJSON, artifactsJSON, memoryJSON, scheduleJSON []byte
		var message, errStr sql.NullString
		var nextRunAt, completedAt sql.NullTime

		err := rows.Scan(
			&task.ID,
			&task.UserID,
			&task.Query,
			&task.Status,
			&planJSON,
			&subTasksJSON,
			&artifactsJSON,
			&memoryJSON,
			&task.Progress,
			&message,
			&errStr,
			&scheduleJSON,
			&nextRunAt,
			&task.RunCount,
			&task.TotalCost,
			&task.CreatedAt,
			&task.UpdatedAt,
			&completedAt,
		)
		if err != nil {
			continue
		}

		if len(planJSON) > 0 {
			json.Unmarshal(planJSON, &task.Plan)
		}
		if len(subTasksJSON) > 0 {
			json.Unmarshal(subTasksJSON, &task.SubTasks)
		}
		if len(artifactsJSON) > 0 {
			json.Unmarshal(artifactsJSON, &task.Artifacts)
		}
		if len(memoryJSON) > 0 {
			json.Unmarshal(memoryJSON, &task.Memory)
		}
		if len(scheduleJSON) > 0 {
			json.Unmarshal(scheduleJSON, &task.Schedule)
		}

		if message.Valid {
			task.Message = message.String
		}
		if errStr.Valid {
			task.Error = errStr.String
		}
		if nextRunAt.Valid {
			task.NextRunAt = &nextRunAt.Time
		}
		if completedAt.Valid {
			task.CompletedAt = &completedAt.Time
		}

		tasks = append(tasks, task)
	}

	return tasks, nil
}

func (r *ComputerTaskRepo) Delete(ctx context.Context, id string) error {
	query := `DELETE FROM computer_tasks WHERE id = $1`
	_, err := r.db.ExecContext(ctx, query, id)
	return err
}

func (r *ComputerTaskRepo) DeleteOlderThan(ctx context.Context, days int) (int64, error) {
	query := `
		DELETE FROM computer_tasks
		WHERE created_at < NOW() - INTERVAL '%d days'
		AND status IN ('completed', 'failed', 'cancelled')
	`
	result, err := r.db.ExecContext(ctx, fmt.Sprintf(query, days))
	if err != nil {
		return 0, err
	}
	return result.RowsAffected()
}

func (r *ComputerTaskRepo) CountByUser(ctx context.Context, userID string) (int64, error) {
	query := `SELECT COUNT(*) FROM computer_tasks WHERE user_id = $1`
	var count int64
	err := r.db.QueryRowContext(ctx, query, userID).Scan(&count)
	return count, err
}

func (r *ComputerTaskRepo) CountByStatus(ctx context.Context, status string) (int64, error) {
	query := `SELECT COUNT(*) FROM computer_tasks WHERE status = $1`
	var count int64
	err := r.db.QueryRowContext(ctx, query, status).Scan(&count)
	return count, err
}