import os
import json
import term
import time

const c = term.colorize
const tg = term.green
const tm = term.magenta
const tb = term.bold

struct Check {
	name     string
	bucket   string
	state    string
	link     string
	workflow string
}

fn (c Check) is_failed() bool {
	return c.bucket == 'fail' || c.state in ['FAILURE', 'TIMED_OUT']
}

fn (c Check) is_cancelled() bool {
	return c.bucket == 'cancel' || c.state == 'CANCELLED'
}

struct GhCheckRun {
	name       string
	status     string
	conclusion string
	html_url   string @[json: html_url]
}

struct GhCheckRunsResponse {
	check_runs []GhCheckRun @[json: check_runs]
}

fn main() {
	unbuffer_stdout()
	arg := os.args[1] or {
		println('Usage: v run gh_restart_failed.v <PR_NUMBER|REF>')
		return
	}
	is_pr := arg.len < 6 && arg.bytes().all(it.is_digit())
	checks := if is_pr { get_checks_for_pr(arg.int()) } else { get_checks_for_commit(arg) }
	mut failed := []Check{}
	mut cancelled := []Check{}
	mut succeeded := 0
	mut in_progress := 0
	for check in checks {
		if check.is_failed() {
			failed << check
		} else if check.is_cancelled() {
			cancelled << check
		} else if check.bucket == 'pass' || check.state == 'SUCCESS' {
			succeeded++
		} else {
			in_progress++
		}
	}
	mut to_restart := []Check{}
	to_restart << failed
	to_restart << cancelled
	mut restarted_count := 0
	if to_restart.len > 0 {
		println('Found ${to_restart.len} failed or cancelled jobs:')
		for job in to_restart {
			println('- ${job.workflow} / ${job.name} (${job.state})')
		}
		println('\n' + c(tg, 'Do you want to restart these ${to_restart.len} jobs? [y/N]: '))
		if os.input('').to_lower().trim_space() == 'y' {
			println('')
			for job in to_restart {
				run_id, job_id := parse_ids(job.link)
				if run_id == '' || job_id == '' {
					println('Could not parse IDs from link: ${job.link} (Skipping)')
					continue
				}
				print('Restarting ${job.name} (Run: ${run_id}, Job: ${job_id})... ')
				res := execute_with_progress('gh run rerun ${run_id} --job ${job_id}')
				if res.exit_code == 0 {
					restarted_count++
				} else {
					println('  Error: ${res.output.trim_space()}')
				}
			}
		} else {
			println('Aborted restart.')
		}
	} else {
		println('No failed or cancelled jobs found.')
	}
	println('\n' + c(tg, 'Summary:'))
	println('Total jobs:   ${m(checks.len)}')
	println('Failed:       ${m(failed.len)}')
	println('Cancelled:    ${m(cancelled.len)}')
	println('Succeeded:    ${m(succeeded)}')
	println('In Progress:  ${m(in_progress)}')
	println('Restarted:    ${m(restarted_count)}')
}

fn m(n int) string {
	return c(tb, n.str())
}

fn parse_ids(link string) (string, string) {
	parts := link.split('/')
	runs_idx := parts.index('runs')
	job_idx := parts.index('job')
	mut run_id := ''
	mut job_id := ''
	if runs_idx != -1 && runs_idx + 1 < parts.len {
		run_id = parts[runs_idx + 1]
	}
	if job_idx != -1 && job_idx + 1 < parts.len {
		job_id = parts[job_idx + 1]
	}
	return run_id, job_id
}

fn get_checks_for_pr(pr_number int) []Check {
	println(c(tg, 'Fetching checks for PR ${m(pr_number)}...'))
	res := execute_with_progress('gh pr checks ${pr_number} --json name,bucket,state,link,workflow')
	if res.exit_code != 0 {
		println('Error: ${res.output}')
		exit(1)
	}
	return json.decode([]Check, res.output) or { exit(1) }
}

fn get_checks_for_commit(commit string) []Check {
	println(c(tg, 'Fetching checks for ref ${c(tb, commit)}...'))
	res := execute_with_progress('gh api repos/:owner/:repo/commits/${commit}/check-runs?per_page=100')
	if res.exit_code != 0 {
		println('Error: ${res.output}')
		exit(1)
	}
	resp := json.decode(GhCheckRunsResponse, res.output) or { exit(1) }
	mut checks := []Check{}
	for cr in resp.check_runs {
		checks << Check{
			name:     cr.name
			bucket:   if cr.conclusion == 'failure' {
				'fail'
			} else if cr.conclusion == 'cancelled' {
				'cancel'
			} else if cr.conclusion == 'success' {
				'pass'
			} else {
				'pending'
			}
			state:    if cr.conclusion != '' {
				cr.conclusion.to_upper()
			} else {
				cr.status.to_upper()
			}
			link:     cr.html_url
			workflow: 'Actions'
		}
	}
	return checks
}

fn execute_with_progress(cmd string) os.Result {
	start := time.now()
	mut stop := false
	spawn fn (cmd string, start time.Time, mut stop &bool) {
		start_str := start.hhmmss()
		for !*stop {
			elapsed := time.since(start).seconds()
			print('\rRunning ${c(tm, cmd)} [${start_str}] ... ${elapsed:.1f}s')
			os.flush()
			time.sleep(100 * time.millisecond)
		}
	}(cmd, start, mut &stop)
	res := os.execute(cmd)
	elapsed := time.since(start).seconds()
	stop = true
	time.sleep(100 * time.millisecond)
	status := if res.exit_code == 0 { 'OK' } else { 'Failed' }
	print('\r')
	println('Command ${c(tm, cmd)} done in ${elapsed:.1f}s. ${status}')
	return res
}