import { LitElement, html, css, isServer, type PropertyValues } from 'lit';
import { customElement, property, state } from 'lit/decorators.js';
import { type ChallengeSolveRequest, type SolutionMessage, WorkerMessageType, type WorkerRequest, WorkerResponseType } from './types/worker';
import { type Challenge, ChallengeStrategy, ChallengeAlgorithm } from '@impost/lib';
import { get_wasm_module } from '@impost/lib/solver';

import ChallengeWorker from './solver-worker?worker&inline';

@customElement('pow-captcha')
export class PowCaptcha extends LitElement {
    static override styles = css`
    :host {
        display: block;
        width: var(--impost-widget-width, 250px);
    }
    .impost-widget {
        display: flex;
        flex-direction: column;
        border: 1px solid var(--impost-widget-border-color, #505050);
        font-family: sans-serif;
        background-color: var(--impost-widget-background-color, #070408);
        border-radius: var(--impost-widget-border-radius, 8px);
    }
    .impost-error-icon {
        color: var(--impost-widget-error-icon-color, #FF8117);
    }
    .impost-main {
        display: flex;
        align-items: center;
        gap: 0.375rem;
        padding: 0.75rem;
    }
    .impost-main span {
        height: min-content;
    }
    .impost-checkbox {
        display: flex;
        align-items: center;
        justify-content: center;
        width: 1.5rem;
        height: 1.5rem;
        color: var(--impost-widget-checkbox-color, #6e6e6eff);
    }
    .impost-checkbox input {
        height: 80%;
        width: 80%;
        margin: 0;
    }
    .impost-footer {
        display: flex;
        padding: 0.25rem 0.75rem;
        font-size: 0.75rem;
        color: var(--impost-widget-footer-color, #a2a2a2);
        justify-content: space-around;
        align-items: end;
    }
    .impost-footer div {
        height: min-content;
    }
    .impost-footer #provider-link {
        margin-left: auto;
    }
    .impost-footer div a {
        color: var(--impost-widget-footer-color, #a2a2a2);
    }
    .impost-footer div a:hover {
        text-decoration: none;
    }
    .hidden {
        display: none;
    }
  `;


    // one of: "load", "focus", "submit", "off"
    @property({ type: String })
    auto: "onload" | "onfocus" | "onsubmit" | "off" = "off";

    @property({ type: String })
    challengeUrl: string = 'http://localhost:3000/api/pow';

    @property({ type: String })
    challengejson: string = '';

    // needed to allow for multiple widgets on the same page if you wanted to
    // do that for some reason
    uid: string = Math.floor(Math.random() * 100000).toString();

    @state()
    private solution: string = '';

    @state()
    private challengeData: Challenge | null = null;

    @state()
    private status: 'unsolved' | 'solving' | 'solved' | 'error' = 'unsolved';

    @state()
    private disabled: boolean = true;

    // stores the nonce and solution atomics
    private sab: SharedArrayBuffer = new SharedArrayBuffer(8);

    private solverWorkers: Worker[] | null = null;

    override connectedCallback() {
        super.connectedCallback();
        this.fetchChallenge();

        this.initWorkers();

        switch (this.auto) {
            case 'onload':
                this.solveChallenge();
                break;
            case 'onfocus':
                this.addEventListener('focus', () => this.solveChallenge());
                break;
            case 'onsubmit':
                if (this.parentElement?.nodeName === 'FORM') {
                    this.parentElement.addEventListener('submit', () => this.solveChallenge());
                }
                break;
        }
    }

    protected override firstUpdated(_changedProperties: PropertyValues): void {
        super.firstUpdated(_changedProperties);

        this.disabled = false;
    }

    override disconnectedCallback() {
        super.disconnectedCallback();

        for (const worker of this.solverWorkers || []) {
            worker.terminate();
            this.solverWorkers = null;
        }
    }

    getCurrentWorkingNonce() {
        return Atomics.load(new Uint32Array(this.sab), 0);
    }

    async fetchChallenge() {
        if (this.challengeData !== null) {
            return;
        }

        if (this.challengejson !== '' && this.challengeData === null) {
            this.challengeData = JSON.parse(this.challengejson);
            this.challengejson = '';
            return;
        }

        if (isServer) {
            return;
        }

        await fetch(`${this.challengeUrl}/challenge`)
            .then(response => response.json())
            .then(data => {
                this.challengeData = data;
            })
            .catch(error => {
                console.error('Error fetching challenge:', error);
                console.error('Failed to fetch challenge');
            });
    }

    async initWorkers() {
        this.solverWorkers = [];

        const num_workers = navigator.hardwareConcurrency || 4;
        for (let i = 0; i < num_workers; i++) {
            this.solverWorkers.push(new ChallengeWorker());
        }

        const atomics_view = new Int32Array(this.sab);
        Atomics.store(atomics_view, 0, 0);
        Atomics.store(atomics_view, 1, 0);

        let wasm_module = await get_wasm_module();
        let worker_promises: Promise<void>[] = [];
        for (let i = 0; i < this.solverWorkers.length; i++) {
            console.log('Worker', i);
            const worker = this.solverWorkers[i]!;

            worker_promises.push(new Promise<void>((resolve, reject) => {
                const message_handler = (event: MessageEvent<SolutionMessage>) => {
                    console.log('Worker', i, 'got message', event.data);

                    if (event.data.type === WorkerResponseType.Error) {
                        console.error("Worker error:", event.data.error);
                        reject(event.data.error);
                    }

                    if (event.data.type === WorkerResponseType.InitOk) {
                        resolve();
                    }

                    reject(new Error("Unexpected message from worker"));
                };

                const error_handler = (error: ErrorEvent) => {
                    console.error("Worker error:", error);
                    reject(error);
                };

                worker.addEventListener('message', message_handler);
                worker.addEventListener('error', error_handler);

                worker.postMessage({
                    type: WorkerMessageType.Init,
                    module: wasm_module,
                    sab: this.sab,
                } as WorkerRequest);
            }));
        }

        const timeoutMs = 10 * 1000;
        let timeout: number;
        const timeoutPromise = new Promise((_, reject) => {
            timeout = setTimeout(() => {
                console.error('Failed to initialize workers in time');
                this.status = 'error';
                reject(new Error(`Function timed out after ${timeoutMs}ms`));
            }, timeoutMs);
        });

        await Promise.race([
            Promise.allSettled(worker_promises).then(() => {
                console.log('All workers initialized');
            }),
            timeoutPromise,
        ]).then(() => {
            clearTimeout(timeout);
        });
    }

    async issueChallengeToWorker(worker: Worker, request: ChallengeSolveRequest): Promise<SolutionMessage> {
        return new Promise<SolutionMessage>((resolve, reject) => {
            const message_handler = (event: MessageEvent<SolutionMessage>) => {
                worker.removeEventListener('message', message_handler);
                worker.removeEventListener('error', error_handler);

                resolve(event.data);
            };

            const error_handler = (error: ErrorEvent) => {
                worker.removeEventListener('error', error_handler);
                worker.removeEventListener('message', message_handler);
                console.error("Worker error:", error);

                reject(error);
            };

            worker.addEventListener('message', message_handler);
            worker.addEventListener('error', error_handler);

            // switch (request.strategy) {
            //     case ChallengeStrategy.LeadingZeroes:
            //         worker.postMessage({
            //             strategy: ChallengeStrategy.LeadingZeroes,
            //             salt: request.salt,
            //             difficulty: request.difficulty,
            //         } as WorkerRequest);
            //         break;
            //     case ChallengeStrategy.TargetNumber:
            //         worker.postMessage({
            //             strategy: ChallengeStrategy.TargetNumber,
            //             target: request.target,
            //             salt: request.salt,
            //         } as WorkerRequest);
            //         break;
            // }
            switch (request.algorithm) {
                case ChallengeAlgorithm.SHA256:
                    switch (request.strategy) {
                        case ChallengeStrategy.LeadingZeroes:
                            worker.postMessage({
                                algorithm: ChallengeAlgorithm.SHA256,
                                strategy: ChallengeStrategy.LeadingZeroes,
                                salt: request.salt,
                                difficulty: request.difficulty,
                            } as WorkerRequest);
                            break;
                        case ChallengeStrategy.TargetNumber:
                            worker.postMessage({
                                algorithm: ChallengeAlgorithm.SHA256,
                                strategy: ChallengeStrategy.TargetNumber,
                                target: request.target,
                                salt: request.salt,
                            } as WorkerRequest);
                            break;
                    }
                    break;
                case ChallengeAlgorithm.Argon2id:
                    switch (request.strategy) {
                        case ChallengeStrategy.LeadingZeroes:
                            worker.postMessage({
                                algorithm: ChallengeAlgorithm.Argon2id,
                                strategy: ChallengeStrategy.LeadingZeroes,
                                salt: request.salt,
                                difficulty: request.difficulty,
                            } as WorkerRequest);
                            break;
                        case ChallengeStrategy.TargetNumber:
                            worker.postMessage({
                                algorithm: ChallengeAlgorithm.Argon2id,
                                strategy: ChallengeStrategy.TargetNumber,
                                target: request.target,
                                salt: request.salt,
                            } as WorkerRequest);
                            break;
                    }
                    break;
                case ChallengeAlgorithm.kCTF:
                    worker.postMessage({
                        algorithm: ChallengeAlgorithm.kCTF,
                        salt: request.salt,
                        difficulty: request.difficulty,
                    } as WorkerRequest);
                    break;
            }
        });
    }

    async solveChallenge() {
        if (!this.challengeData || this.solverWorkers === null) {
            console.error('solveChallenge called before challenge is ready');
            this.status = 'error';
            return;
        }

        if (this.solution !== '') {
            return;
        }

        this.dispatchEvent(new CustomEvent('impost:solve', {
            bubbles: true,
            composed: true,
        }))

        console.log(this.challengeData);

        this.status = 'solving';
        this.solution = '';

        const atomics_view = new Int32Array(this.sab);
        Atomics.store(atomics_view, 0, 0);
        Atomics.store(atomics_view, 1, -1);

        let request: ChallengeSolveRequest;

        // switch (this.challengeData.strategy) {
        //     case ChallengeStrategy.LeadingZeroes:
        //         request = {
        //             strategy: ChallengeStrategy.LeadingZeroes,
        //             salt: this.challengeData.salt,
        //             difficulty: this.challengeData.difficulty,
        //         };
        //         break;
        //     case ChallengeStrategy.TargetNumber:
        //         request = {
        //             strategy: ChallengeStrategy.TargetNumber,
        //             target: this.challengeData.target,
        //             salt: this.challengeData.salt,
        //         };
        //         break;
        // }
        switch (this.challengeData.algorithm) {
            case ChallengeAlgorithm.SHA256:
                switch (this.challengeData.strategy) {
                    case ChallengeStrategy.LeadingZeroes:
                        request = {
                            algorithm: ChallengeAlgorithm.SHA256,
                            strategy: ChallengeStrategy.LeadingZeroes,
                            salt: this.challengeData.salt,
                            difficulty: this.challengeData.difficulty,
                        };
                        break;
                    case ChallengeStrategy.TargetNumber:
                        request = {
                            algorithm: ChallengeAlgorithm.SHA256,
                            strategy: ChallengeStrategy.TargetNumber,
                            target: this.challengeData.target,
                            salt: this.challengeData.salt,
                        };
                        break;
                }
                break;
            case ChallengeAlgorithm.Argon2id:
                switch (this.challengeData.strategy) {
                    case ChallengeStrategy.LeadingZeroes:
                        request = {
                            algorithm: ChallengeAlgorithm.Argon2id,
                            strategy: ChallengeStrategy.LeadingZeroes,
                            salt: this.challengeData.salt,
                            difficulty: this.challengeData.difficulty,
                        };
                        break;
                    case ChallengeStrategy.TargetNumber:
                        request = {
                            algorithm: ChallengeAlgorithm.Argon2id,
                            strategy: ChallengeStrategy.TargetNumber,
                            target: this.challengeData.target,
                            salt: this.challengeData.salt,
                        };
                        break;
                }
                break;
            case ChallengeAlgorithm.kCTF:
                request = {
                    algorithm: ChallengeAlgorithm.kCTF,
                    salt: this.challengeData.salt,
                    difficulty: this.challengeData.difficulty,
                };
                break;
        }

        console.log('Sending challenge to workers...');
        // TODO: the first response is not always the solution, due to cmpxchg 
        // blocking, some workers may block on the read, and as soon as they 
        // unblock, they return 0 since the challenge is already solved.
        //
        // TODO: We need to do a better job of tracking solvers, so if one worker
        // errors out, we only error out if all workers have errored out.
        let worker_promises: Promise<SolutionMessage>[] = [];
        if (request.algorithm === ChallengeAlgorithm.kCTF) {
            worker_promises.push(this.issueChallengeToWorker(this.solverWorkers[0], request));
        } else {
            for (let worker of this.solverWorkers) {
                // dispatch to all workers, func is async so it will not block
                worker_promises.push(this.issueChallengeToWorker(worker, request));
            }
        }

        let solution = await Promise.race(worker_promises);

        if (solution.type === WorkerResponseType.Error) {
            console.error("Worker error:", solution.error);
            this.status = 'error';
            return;
        }

        if (solution.type !== WorkerResponseType.Solution) {
            console.error("Worker sent spurious message");
            this.status = 'error';
            return;
        }

        // TODO: configure if we should fetch or not
        try {
            await fetch(`${this.challengeUrl}/challenge`, {
                method: 'POST',
                body: JSON.stringify({
                    salt: this.challengeData.salt,
                    solution: solution.solution,
                }),
                headers: {
                    'Content-Type': 'application/json'
                }
            })

            this.status = 'solved';

            this.dispatchEvent(new CustomEvent('impost:solved', {
                detail: {
                    salt: this.challengeData.salt,
                    solution: solution.solution,
                },
                bubbles: true,
                composed: true,
            }))
        } catch {
            console.error('Failed to submit solution');
            this.status = 'error';
        }
    }

    solvePreventDefault(event: Event) {
        event.preventDefault();

        this.solveChallenge();
    }

    override render() {
        if (this.challengejson !== '' && this.challengeData === null) {
            this.challengeData = JSON.parse(this.challengejson);
            this.challengejson = '';
        }

        if (this.challengeData === null) {
            return html`
        <div class="loading-message">Loading captcha challenge...</div>
      `;
        }

        return html`
        <div class="impost-widget">
            <div class="impost-main">
                <div class="impost-checkbox">
                    ${this.status !== 'solving' ? html`${this.status === 'error' ? html`<svg class="impost-error-icon" xmlns="http://www.w3.org/2000/svg" width="32" height="32" viewBox="0 0 24 24"><!-- Icon from Material Symbols by Google - https://github.com/google/material-design-icons/blob/master/LICENSE --><path fill="currentColor" d="M12 14q-.425 0-.712-.288T11 13V6q0-.425.288-.712T12 5t.713.288T13 6v7q0 .425-.288.713T12 14m0 5q-.425 0-.712-.288T11 18t.288-.712T12 17t.713.288T13 18t-.288.713T12 19"/></svg>` : html`
                        <input type="checkbox" id="impost-checkbox-${this.uid}" @click=${this.solvePreventDefault} @change=${this.solvePreventDefault} ?disabled=${this.disabled} ?checked=${this.status === 'solved'}>
                    `}` : html`
                        <svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24">
                            <!-- Icon from SVG Spinners by Utkarsh Verma - https://github.com/n3r4zzurr0/svg-spinners/blob/main/LICENSE -->
                            <g stroke="currentColor">
                                <circle cx="12" cy="12" r="9.5" fill="none" stroke-linecap="round" stroke-width="3">
                                    <animate attributeName="stroke-dasharray" calcMode="spline" dur="1.5s" keySplines="0.42,0,0.58,1;0.42,0,0.58,1;0.42,0,0.58,1" keyTimes="0;0.475;0.95;1" repeatCount="indefinite" values="0 150;42 150;42 150;42 150"/>
                                    <animate attributeName="stroke-dashoffset" calcMode="spline" dur="1.5s" keySplines="0.42,0,0.58,1;0.42,0,0.58,1;0.42,0,0.58,1" keyTimes="0;0.475;0.95;1" repeatCount="indefinite" values="0;-16;-59;-59"/>
                                </circle>
                                <animateTransform attributeName="transform" dur="2s" repeatCount="indefinite" type="rotate" values="0 12 12;360 12 12"/>
                            </g>
                        </svg>
                    `}
                </div>
                <label for="impost-checkbox-${this.uid}">${this.status === 'error' ? 'Something went wrong' : this.status === 'solved' ? 'Verified' : html`${this.status === 'solving' ? 'Verifying...' : 'I am not a robot'}`}</label>
            </div>

            <div class="impost-footer">
                <div id="provider-link">
                    Protected by <a href="https://github.com/impost/pow-captcha" target="_blank">Impost</a>
                </div>
            </div>
        </div>

        <input
            type = "text"
            id = "impost-solution"
            class="hidden"
            .value = ${this.solution}
        />
    `;
    }
}