import argparse
import os
import threading
import time
import glfw
from OpenGL.GL import *

import mersenne_sieve

SHADER_PATH = os.path.join(os.path.dirname(__file__), 'mersenne_field.comp')
WINDOW_SIZE = 1024

CANDIDATES = [22514117, 22510841, 22514533]
verification_state = {
    'stage': 'idle',
    'progress_trial': [0.0, 0.0, 0.0],
    'progress_ll': [0.0, 0.0, 0.0],
    'feature_low': [0.0, 0.0, 0.0],
    'feature_high': [0.0, 0.0, 0.0],
    'feature_mid': [0.0, 0.0, 0.0],
    'hdgl_score': [0.0, 0.0, 0.0],
    'exponent_prime': [0.0, 0.0, 0.0],
    'trial_factor_found': [0.0, 0.0, 0.0],
    'current_idx': 0,
    'message': 'waiting'
}
state_lock = threading.Lock()

VERT_SHADER = '''#version 330 core
out vec2 vUV;
void main() {
    vec2 positions[3] = vec2[3](
        vec2(-1.0, -1.0),
        vec2(3.0, -1.0),
        vec2(-1.0, 3.0)
    );
    vec2 pos = positions[gl_VertexID];
    vUV = pos * 0.5 + 0.5;
    gl_Position = vec4(pos, 0.0, 1.0);
}
'''

FRAG_SHADER = '''#version 330 core
in vec2 vUV;
uniform sampler2D uTexture;
out vec4 FragColor;
void main() {
    FragColor = texture(uTexture, vUV);
}
'''


def compile_shader(source, shader_type):
    shader = glCreateShader(shader_type)
    glShaderSource(shader, source)
    glCompileShader(shader)
    status = glGetShaderiv(shader, GL_COMPILE_STATUS)
    if not status:
        error = glGetShaderInfoLog(shader).decode('utf-8')
        raise RuntimeError(f'Shader compile error: {error}')
    return shader


def link_program(*shaders):
    program = glCreateProgram()
    for shader in shaders:
        glAttachShader(program, shader)
    glLinkProgram(program)
    status = glGetProgramiv(program, GL_LINK_STATUS)
    if not status:
        error = glGetProgramInfoLog(program).decode('utf-8')
        raise RuntimeError(f'Program link error: {error}')
    for shader in shaders:
        glDetachShader(program, shader)
        glDeleteShader(shader)
    return program


def load_compute_shader(path):
    with open(path, 'r', encoding='utf-8') as f:
        source = f.read()
    return compile_shader(source, GL_COMPUTE_SHADER)


def build_compute_program():
    shader = load_compute_shader(SHADER_PATH)
    return link_program(shader)


def build_screen_program():
    vert = compile_shader(VERT_SHADER, GL_VERTEX_SHADER)
    frag = compile_shader(FRAG_SHADER, GL_FRAGMENT_SHADER)
    return link_program(vert, frag)


def make_texture(size):
    tex = glGenTextures(1)
    glBindTexture(GL_TEXTURE_2D, tex)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, size, size, 0, GL_RGBA, GL_FLOAT, None)
    glBindTexture(GL_TEXTURE_2D, 0)
    return tex


def main():
    if not glfw.init():
        raise RuntimeError('glfw initialization failed')

    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 4)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    glfw.window_hint(glfw.RESIZABLE, glfw.FALSE)

    window = glfw.create_window(WINDOW_SIZE, WINDOW_SIZE, 'Mersenne Field Visualizer', None, None)
    if not window:
        glfw.terminate()
        raise RuntimeError('glfw window creation failed')

    glfw.make_context_current(window)

    compute_program = build_compute_program()
    screen_program = build_screen_program()
    texture = make_texture(WINDOW_SIZE)
    vao = glGenVertexArrays(1)
    glBindVertexArray(vao)

    start_time = time.time()

    progress_loc = glGetUniformLocation(compute_program, 'progressTrial')
    ll_loc = glGetUniformLocation(compute_program, 'progressLL')
    feature_low_loc = glGetUniformLocation(compute_program, 'featureLow')
    feature_high_loc = glGetUniformLocation(compute_program, 'featureHigh')
    feature_mid_loc = glGetUniformLocation(compute_program, 'featureMid')
    score_loc = glGetUniformLocation(compute_program, 'hdglScore')
    prime_loc = glGetUniformLocation(compute_program, 'exponentPrime')
    factor_loc = glGetUniformLocation(compute_program, 'trialFactorFound')
    exponent_loc = glGetUniformLocation(compute_program, 'candidateExp')
    active_loc = glGetUniformLocation(compute_program, 'progressActive')
    time_loc = glGetUniformLocation(compute_program, 'time')

    def state_features_from_s(s=None, p=None):
        if s is None or p is None:
            if p is None:
                return 0.0, 0.0, 0.0
            base = float(p % 100000) / 100000.0
            mid = float((p // 100000) % 100000) / 100000.0
            high = float((p // 10000000000) % 100000) / 100000.0
            return base, mid, high
        mask = (1 << 24) - 1
        low = float(int(s & mask)) / float(mask)
        high = float(int((s >> max(0, p - 24)) & mask)) / float(mask)
        mid = float(int((s >> max(0, p // 2)) & mask)) / float(mask)
        return low, high, mid

    def progress_callback(stage, progress, candidate_idx, s=None, p=None):
        with state_lock:
            if stage == 'trial':
                verification_state['progress_trial'][candidate_idx] = progress
                verification_state['message'] = f'Candidate {candidate_idx + 1} trial {int(progress * 100)}%'
            elif stage == 'll':
                verification_state['progress_ll'][candidate_idx] = progress
                verification_state['message'] = f'Candidate {candidate_idx + 1} LL {int(progress * 100)}%'
            else:
                verification_state['message'] = str(stage)
            low, high, mid = state_features_from_s(s, p)
            verification_state['feature_low'][candidate_idx] = low
            verification_state['feature_high'][candidate_idx] = high
            verification_state['feature_mid'][candidate_idx] = mid

    def verification_worker():
        for idx, p in enumerate(CANDIDATES):
            with state_lock:
                verification_state['current_idx'] = idx
                verification_state['stage'] = 'verify'
                verification_state['message'] = f'Verify p={p}'
                verification_state['progress_trial'][idx] = 0.0
                verification_state['progress_ll'][idx] = 0.0
            def callback(stage, progress, s=None, p=None, idx_inner=idx):
                progress_callback(stage, progress, idx_inner, s=s, p=p)
            result = mersenne_sieve.verify_candidate(p, progress_callback=callback)
            with state_lock:
                verification_state['hdgl_score'][idx] = float(result.get('hdgl_score', 0.0) or 0.0)
                verification_state['exponent_prime'][idx] = 1.0 if result.get('prime_exponent') else 0.0
                verification_state['trial_factor_found'][idx] = 1.0 if result.get('trial_factor') else 0.0
                verification_state['stage'] = 'done'
                verification_state['message'] = f'Candidate {idx + 1} done'

    parser = argparse.ArgumentParser(description='Mersenne field visualizer with live verification progress')
    parser.add_argument('--duration', type=float, default=None, help='Run duration in seconds before exiting')
    args = parser.parse_args()

    verifier_thread = threading.Thread(target=verification_worker, daemon=True)
    verifier_thread.start()
    end_time = time.time() + args.duration if args.duration is not None else None

    glUseProgram(screen_program)
    glUniform1i(glGetUniformLocation(screen_program, 'uTexture'), 0)

    while not glfw.window_should_close(window):
        if end_time is not None and time.time() >= end_time:
            break
        current_time = time.time() - start_time

        with state_lock:
            pt = verification_state['progress_trial']
            pll = verification_state['progress_ll']
            low = verification_state['feature_low']
            high = verification_state['feature_high']
            mid = verification_state['feature_mid']
            score = verification_state['hdgl_score']
            prime = verification_state['exponent_prime']
            factor = verification_state['trial_factor_found']
            exponent_vals = [float(x) for x in CANDIDATES]
            active = 1.0 if verification_state['stage'] != 'idle' else 0.0
            message = verification_state['message']
            current_idx = verification_state['current_idx']

        glUseProgram(compute_program)
        glBindImageTexture(0, texture, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA32F)
        glUniform1f(time_loc, float(current_time))
        if progress_loc != -1:
            glUniform1fv(progress_loc, len(pt), pt)
        if ll_loc != -1:
            glUniform1fv(ll_loc, len(pll), pll)
        if feature_low_loc != -1:
            glUniform1fv(feature_low_loc, len(low), low)
        if feature_high_loc != -1:
            glUniform1fv(feature_high_loc, len(high), high)
        if feature_mid_loc != -1:
            glUniform1fv(feature_mid_loc, len(mid), mid)
        if score_loc != -1:
            glUniform1fv(score_loc, len(score), score)
        if prime_loc != -1:
            glUniform1fv(prime_loc, len(prime), prime)
        if factor_loc != -1:
            glUniform1fv(factor_loc, len(factor), factor)
        if exponent_loc != -1:
            glUniform1fv(exponent_loc, len(exponent_vals), exponent_vals)
        if active_loc != -1:
            glUniform1f(active_loc, float(active))
        glDispatchCompute(WINDOW_SIZE // 16, WINDOW_SIZE // 16, 1)
        glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT)
        glfw.set_window_title(window, f'Mersenne Field Visualizer - candidate {current_idx + 1} {message}')

        glClear(GL_COLOR_BUFFER_BIT)
        glUseProgram(screen_program)
        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, texture)
        glDrawArrays(GL_TRIANGLES, 0, 3)

        glfw.swap_buffers(window)
        glfw.poll_events()

    glDeleteTextures([texture])
    glDeleteProgram(compute_program)
    glDeleteProgram(screen_program)
    glfw.terminate()


if __name__ == '__main__':
    main()
