package storage

import (
	"context"
	"database/sql"
	"fmt"
)

// AppendTrace добавляет трассу выполнения к задаче. Возвращает её ID.
func (s *Storage) AppendTrace(ctx context.Context, tr *Trace) (int64, error) {
	now := Now()
	res, err := s.db.ExecContext(ctx, `
		INSERT INTO traces (task_id, agent, session_id, prompt, output, status, started_at)
		VALUES (?, ?, ?, ?, ?, ?, ?)`,
		tr.TaskID, tr.Agent, tr.SessionID, tr.Prompt, tr.Output, TraceRunning, now,
	)
	if err != nil {
		return 0, fmt.Errorf("%w: append trace: %w", ErrDB, err)
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("%w: last insert id: %w", ErrDB, err)
	}
	tr.ID = id
	tr.Status = TraceRunning
	tr.StartedAt = now
	return id, nil
}

// UpdateTraceStatus обновляет статус и finished_at трассы.
func (s *Storage) UpdateTraceStatus(ctx context.Context, traceID int64, status TraceStatus) error {
	now := Now()
	res, err := s.db.ExecContext(ctx, `
		UPDATE traces SET status=?, finished_at=? WHERE id=?`,
		status, now, traceID,
	)
	if err != nil {
		return fmt.Errorf("%w: update trace %d: %w", ErrDB, traceID, err)
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return fmt.Errorf("%w: trace %d", ErrTraceNotFound, traceID)
	}
	return nil
}

// UpdateTraceOutput обновляет output трассы (например, по мере поступления данных).
func (s *Storage) UpdateTraceOutput(ctx context.Context, traceID int64, output string) error {
	_, err := s.db.ExecContext(ctx, `UPDATE traces SET output=? WHERE id=?`, output, traceID)
	if err != nil {
		return fmt.Errorf("%w: update trace output %d: %w", ErrDB, traceID, err)
	}
	return nil
}

// GetTraces возвращает трассы для задачи, отсортированные по started_at.
func (s *Storage) GetTraces(ctx context.Context, taskID int64) ([]*Trace, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT id, task_id, agent, session_id, prompt, output, status, started_at, finished_at
		FROM traces WHERE task_id=? ORDER BY started_at ASC`, taskID)
	if err != nil {
		return nil, fmt.Errorf("%w: get traces for task %d: %w", ErrDB, taskID, err)
	}
	defer rows.Close()

	var traces []*Trace
	for rows.Next() {
		tr := &Trace{}
		var finishedAt NullSQLiteTime
		if err := rows.Scan(&tr.ID, &tr.TaskID, &tr.Agent, &tr.SessionID,
			&tr.Prompt, &tr.Output, &tr.Status, &tr.StartedAt, &finishedAt); err != nil {
			return nil, fmt.Errorf("%w: scan trace: %w", ErrDB, err)
		}
		tr.FinishedAt = finishedAt
		traces = append(traces, tr)
	}
	return traces, rows.Err()
}

// GetLatestTrace возвращает последнюю трассу задачи для заданного агента.
func (s *Storage) GetLatestTrace(ctx context.Context, taskID int64, agent string) (*Trace, error) {
	tr := &Trace{}
	var finishedAt NullSQLiteTime
	err := s.db.QueryRowContext(ctx, `
		SELECT id, task_id, agent, session_id, prompt, output, status, started_at, finished_at
		FROM traces WHERE task_id=? AND agent=? ORDER BY id DESC LIMIT 1`,
		taskID, agent).Scan(
		&tr.ID, &tr.TaskID, &tr.Agent, &tr.SessionID,
		&tr.Prompt, &tr.Output, &tr.Status, &tr.StartedAt, &finishedAt,
	)
	if err == sql.ErrNoRows {
		return nil, fmt.Errorf("%w: trace for task %d agent %s", ErrTraceNotFound, taskID, agent)
	}
	if err != nil {
		return nil, fmt.Errorf("%w: get latest trace: %w", ErrDB, err)
	}
	tr.FinishedAt = finishedAt
	return tr, nil
}

// DeleteTrace удаляет трассу. Только для тестов/админки.
func (s *Storage) DeleteTrace(ctx context.Context, id int64) error {
	res, err := s.db.ExecContext(ctx, `DELETE FROM traces WHERE id = ?`, id)
	if err != nil {
		return fmt.Errorf("%w: delete trace %d: %w", ErrDB, id, err)
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return fmt.Errorf("%w: trace %d", ErrTraceNotFound, id)
	}
	return nil
}