import { AsyncExecutionResult, VertesiaClient } from "@vertesia/client";
import { AgentSearchScope, ConversationVisibility, ExecutionEnvironmentRef, InCodeInteraction, JSONSchema, mergeInCodePromptSchemas, supportsToolUse, UserChannel, WorkflowInteractionVars } from "@vertesia/common";
import { JSONObject } from "@vertesia/json";
import { useUserSession } from "@vertesia/ui/session";
import Ajv, { ValidateFunction } from "ajv";
import React, { createContext, useContext, useState, useSyncExternalStore } from "react";

export type WorkflowMode = 'start' | 'schedule';
type ModelOptions = NonNullable<WorkflowInteractionVars["config"]["model_options"]>;

export interface ScheduledWorkflowConfig {
    name: string;
    description?: string;
    cron_expression: string;
    timezone: string;
}

export class PayloadBuilderStore {
    private _listeners = new Set<() => void>();
    snapshot: PayloadBuilder;

    constructor(client: VertesiaClient) {
        this.snapshot = new PayloadBuilder(client, this);
    }

    subscribe = (listener: () => void) => {
        this._listeners.add(listener);
        return () => { this._listeners.delete(listener); };
    };

    getSnapshot = () => this.snapshot;

    notify() {
        this._listeners.forEach(listener => listener());
    }
}

export class PayloadBuilder {
    _interactive: boolean = true;
    _debug_mode: boolean = false;
    _non_blocking_subagents: boolean = true;
    _checkpoint_tokens: number | undefined;
    _visibility: ConversationVisibility | undefined;
    _user_channels: UserChannel[] | undefined;
    _collection: string | undefined;
    _start: boolean = false;
    _preserveRunValues: boolean = false;
    _interaction: InCodeInteraction | undefined;
    _environment: ExecutionEnvironmentRef | undefined;
    _model: string = '';
    _model_options: ModelOptions | undefined;
    _tool_names: string[] = [];
    _data: JSONObject | undefined;
    _mode: WorkflowMode = 'start';
    _scheduledWorkflowConfig: ScheduledWorkflowConfig | undefined;

    private _interactionParamsSchema?: JSONSchema | null;
    private _schemaVersion = 0;
    private _inputValidator?: {
        validate: ValidateFunction;
        schema: JSONSchema;
    };

    private _store: PayloadBuilderStore;

    constructor(public vertesia: VertesiaClient, store: PayloadBuilderStore) {
        this._store = store;
    }

    onStateChanged() {
        const newInstance = this.clone();
        this._store.snapshot = newInstance;
        this._store.notify();
    }

    clone() {
        const builder = new PayloadBuilder(this.vertesia, this._store);
        builder._interactionParamsSchema = this._interactionParamsSchema;
        builder._schemaVersion = this._schemaVersion;
        builder._interaction = this._interaction;
        builder._data = this._data;
        builder._environment = this._environment;
        builder._model = this._model;
        builder._model_options = this._model_options ? { ...this._model_options } as ModelOptions : undefined;
        builder._tool_names = [...this._tool_names];
        builder._interactive = this._interactive;
        builder._debug_mode = this._debug_mode;
        builder._non_blocking_subagents = this._non_blocking_subagents;
        builder._checkpoint_tokens = this._checkpoint_tokens;
        builder._visibility = this._visibility;
        builder._user_channels = this._user_channels ? [...this._user_channels] : undefined;
        builder._inputValidator = this._inputValidator;
        builder._start = this._start;
        builder._collection = this._collection;
        builder._preserveRunValues = this._preserveRunValues;
        builder._mode = this._mode;
        builder._scheduledWorkflowConfig = this._scheduledWorkflowConfig;
        return builder;
    }

    set mode(mode: 'start' | 'schedule') {
        if (mode !== this._mode) {
            this._mode = mode;
            if (mode === 'schedule' && !this._scheduledWorkflowConfig) {
                this._scheduledWorkflowConfig = {
                    name: '',
                    description: '',
                    cron_expression: '0 9 * * *',
                    timezone: 'UTC'
                } as ScheduledWorkflowConfig;
            }
            this.onStateChanged();
        }
    }

    get mode() {
        return this._mode;
    }

    set scheduledWorkflowConfig(config: ScheduledWorkflowConfig | undefined) {
        this._scheduledWorkflowConfig = config;
        this.onStateChanged();
    }

    get scheduledWorkflowConfig() {
        return this._scheduledWorkflowConfig;
    }

    updateScheduledWorkflowConfig(patch: Partial<ScheduledWorkflowConfig>) {
        this._scheduledWorkflowConfig = {
            ...this._scheduledWorkflowConfig,
            ...patch
        } as ScheduledWorkflowConfig;
        this.onStateChanged();
    }

    get interactive() {
        return this._interactive;
    }

    set interactive(interactive: boolean) {
        if (interactive !== this._interactive) {
            this._interactive = interactive;
            this.onStateChanged();
        }
    }

    get debug_mode() {
        return this._debug_mode;
    }

    set debug_mode(debug_mode: boolean) {
        if (debug_mode !== this._debug_mode) {
            this._debug_mode = debug_mode;
            this.onStateChanged();
        }
    }

    get non_blocking_subagents() {
        return this._non_blocking_subagents;
    }

    set non_blocking_subagents(value: boolean) {
        if (value !== this._non_blocking_subagents) {
            this._non_blocking_subagents = value;
            this.onStateChanged();
        }
    }

    get checkpoint_tokens(): number | undefined {
        return this._checkpoint_tokens;
    }

    set checkpoint_tokens(value: number | undefined) {
        if (value !== this._checkpoint_tokens) {
            this._checkpoint_tokens = value;
            this.onStateChanged();
        }
    }

    get visibility(): ConversationVisibility | undefined {
        return this._visibility;
    }

    set visibility(value: ConversationVisibility | undefined) {
        if (value !== this._visibility) {
            this._visibility = value;
            this.onStateChanged();
        }
    }

    get user_channels(): UserChannel[] | undefined {
        return this._user_channels;
    }

    set user_channels(user_channels: UserChannel[] | undefined) {
        this._user_channels = user_channels;
        this.onStateChanged();
    }

    get collection() {
        return this._collection;
    }

    set collection(collection: string | undefined) {
        if (collection !== this._collection) {
            this._collection = collection;
            this.onStateChanged();
        }
    }

    get search_scope() {
        return this._collection ? AgentSearchScope.Collection : undefined;
    }

    async restoreConversation(context: WorkflowInteractionVars) {
        
        // Handle version-specific interaction resolution
        let interactionRef = context.interaction;
        if (context.version) {
            const objectIdRegex = /^[a-fA-F0-9]{24}$/;
            if (!objectIdRegex.test(interactionRef)) {
                // regex to check if interactionRef is an object id (24 hex characters), only append version if not an object id
                interactionRef = `${interactionRef}@${context.version}`;
            }
        }
        
        const inter = await this.vertesia.interactions.catalog.resolve(interactionRef);
        const envId = inter.runtime?.environment || context.config?.environment;
        const model = context.config?.model;
        const env = await (envId ?
            this.vertesia.environments.retrieve(envId).catch(() => undefined)
            :
            Promise.resolve(undefined)
        );


        // Set data BEFORE the interaction setter so that initializeBooleanDefaults (called
        // inside the interactionParamsSchema setter) uses the real run values as its base
        // rather than an empty object. This prevents an intermediate render where the schema
        // is set but _data is still {} — which would cause Input to mount with empty state.
        this._data = context.data;

        // Set interaction (which recomputes interactionParamsSchema from prompts via setter)
        this.interaction = inter;
        // Prefer the schema stored in context (persisted with the run) over the recomputed one
        if (context.interactionParamsSchema != null) {
            this.interactionParamsSchema = context.interactionParamsSchema;
        }

        this._tool_names = context.tool_names || [];
        this._interactive = context.interactive;
        this._debug_mode = context.debug_mode ?? false;
        this._non_blocking_subagents = context.non_blocking_subagents ?? true;
        this._checkpoint_tokens = context.checkpoint_tokens;
        this._user_channels = context.user_channels;
        this._model_options = context.config?.model_options as ModelOptions | undefined;
        this.collection = context.collection_id ?? undefined;

        // we need to trigger the setter to deal with default models
        this.environment = env;
        if (model) {
            this._model = model;
        }

        this.onStateChanged();
    }

    get interaction() {
        return this._interaction;
    }
    set interaction(interaction: InCodeInteraction | undefined) {
        if (interaction?.id !== this._interaction?.id) {
            this._interaction = interaction;
            // trigger the setter to update the onChange state
            this.interactionParamsSchema = interaction ? mergeInCodePromptSchemas(interaction.prompts) : undefined;
            // Reset the validator when schema changes
            this._inputValidator = undefined;
            if (interaction && !this._preserveRunValues) {
                this._model_options = interaction.model_options as ModelOptions | undefined;
                if (interaction.runtime?.environment) {
                    const envId = interaction.runtime.environment;
                    this.vertesia.environments.retrieve(envId).then((environment) => this.environment = environment);
                }
            }
            this.onStateChanged();
        }
    }

    get environment() {
        return this._environment;
    }
    set environment(environment: ExecutionEnvironmentRef | undefined) {
        if (environment?.id !== this._environment?.id) {
            this._environment = environment;
            if (!this._preserveRunValues) {
                // First try to use the interaction model, then the environment default model
                const interactionModel = this.interaction?.runtime?.model;
                if (interactionModel && environment && supportsToolUse(interactionModel, environment.provider)) {
                    this._model = interactionModel;
                } else {
                    this._model = environment?.default_model && supportsToolUse(environment.default_model, environment.provider)
                        ? environment.default_model : '';
                }
            }

            this.onStateChanged();
        }
    }

    get model() {
        return this._model;
    }
    set model(model: string | undefined) {
        if (model !== this._model) {
            this._model = model || '';
            this.onStateChanged();
        }
    }

    get model_options() {
        return this._model_options;
    }

    set model_options(modelOptions: ModelOptions | undefined) {
        this._model_options = modelOptions;
        this.onStateChanged();
    }

    get tool_names() {
        return this._tool_names;
    }
    set tool_names(tools: string[]) {
        this._tool_names = tools;
        this.onStateChanged();
    }

    get data(): JSONObject | undefined {
        return this._data;
    }
    set data(prompt_data: JSONObject) {
        this._data = prompt_data;
        this.onStateChanged();
    }

    set run(run: AsyncExecutionResult | { workflow_id: string; run_id: string }) {
        console.log("run", run);
        this.onStateChanged();
    }

    set start(value: boolean) {
        if (this._start !== value) {
            this._start = value;
            this.onStateChanged();
        }
    }

    markStarted() {
        this.start = true;
    }

    // Method-style setters for use in React components (avoids react-hooks/immutability lint errors)
    setMode(mode: 'start' | 'schedule') { this.mode = mode; }
    setInteraction(interaction: InCodeInteraction | undefined) { this.interaction = interaction; }
    setEnvironment(environment: ExecutionEnvironmentRef | undefined) { this.environment = environment; }
    setModel(model: string | undefined) { this.model = model; }
    setModelOptions(modelOptions: ModelOptions | undefined) { this.model_options = modelOptions; }
    setToolNames(tools: string[]) { this.tool_names = tools; }
    setCollection(collection: string | undefined) { this.collection = collection; }
    setInteractive(interactive: boolean) { this.interactive = interactive; }
    setDebugMode(debug_mode: boolean) { this.debug_mode = debug_mode; }
    setUserChannels(channels: UserChannel[] | undefined) { this.user_channels = channels; }
    setCheckpointTokens(value: number | undefined) { this.checkpoint_tokens = value; }
    setVisibility(value: ConversationVisibility | undefined) { this.visibility = value; }
    setData(data: JSONObject) { this.data = data; }
    setPreserveRunValues(value: boolean) { this.preserveRunValues = value; }

    get start(): boolean {
        return this._start;
    }

    get preserveRunValues(): boolean {
        return this._preserveRunValues;
    }

    set preserveRunValues(value: boolean) {
        this._preserveRunValues = value;
    }

    get schemaVersion(): number {
        return this._schemaVersion;
    }

    get interactionParamsSchema(): JSONSchema | null | undefined {
        return this._interactionParamsSchema;
    }

    set interactionParamsSchema(schema: JSONSchema | null | undefined) {
        if (this._interactionParamsSchema !== schema) {
            this._interactionParamsSchema = schema;
            this._schemaVersion += 1;
            // Booleans must be true or false, never undefined
            if (schema) {
                this._data = this.initializeBooleanDefaults(this._data || {}, schema);
            }
            this.onStateChanged();
        }
    }

    private initializeBooleanDefaults(data: JSONObject, schema: JSONSchema): JSONObject {
        if (!schema.properties) {
            return data;
        }

        const result = { ...data };

        for (const [name, propSchema] of Object.entries(schema.properties)) {
            const prop = propSchema;
            // Initialize boolean fields to false if not already set
            if (prop.type === "boolean" && result[name] === undefined) {
                result[name] = false;
            } else if (prop.type === "object" && prop.properties) {
                // Recursively initialize nested object booleans
                result[name] = this.initializeBooleanDefaults(
                    (result[name] as JSONObject) || {},
                    prop
                );
            }
        }

        return result;
    }

    reset() {
        this._start = false;
        this._interactive = true;
        this._debug_mode = false;
        this._non_blocking_subagents = true;
        this._checkpoint_tokens = undefined;
        this._visibility = undefined;
        this._user_channels = undefined;
        this._collection = undefined;
        this._preserveRunValues = false;
        this._model = '';
        this._model_options = undefined;
        this._environment = undefined;
        this._tool_names = [];
        this._interaction = undefined;
        this._data = undefined;
        this._interactionParamsSchema = null;
        this._inputValidator = undefined;
        this.model = undefined;
        this.environment = undefined;

        this.onStateChanged();

        if (location.hash) {
            const urlWithoutHash = window.location.origin + window.location.pathname + window.location.search;
            history.replaceState(null, '', urlWithoutHash); location.hash = '';
        }
    }

    validateInput(): { isValid: boolean; errorMessage?: string } {
        if (!this._interactionParamsSchema) {
            return { isValid: true };
        }

        // If schema has changed or validator not initialized, recompile
        if (!this._inputValidator || this._inputValidator.schema !== this._interactionParamsSchema) {
            const ajv = new Ajv({ strict: false });
            this._inputValidator = {
                validate: ajv.compile(this._interactionParamsSchema),
                schema: this._interactionParamsSchema
            };
        }

        const prompt_data = this._data || {};
        const isValid = this._inputValidator.validate(prompt_data);

        if (!isValid) {
            const errorMessage: string = this._inputValidator.validate.errors
                ? this._inputValidator.validate.errors.map((err) => `${err.instancePath}: ${err.message}`).join(', ')
                : 'Invalid payload data';
            return { isValid: false, errorMessage };
        }

        return { isValid: true };
    }
}

export const PayloadContext = createContext<PayloadBuilder | undefined>(undefined);

interface PayloadProviderProps {
    children: React.ReactNode;
}
export function PayloadBuilderProvider({ children }: PayloadProviderProps) {
    const { client } = useUserSession();
    const [store] = useState(() => new PayloadBuilderStore(client));
    const builder = useSyncExternalStore(store.subscribe, store.getSnapshot);
    return (
        <PayloadContext.Provider value={builder}>{children}</PayloadContext.Provider >
    )
}

export function usePayloadBuilder() {
    const ctx = useContext(PayloadContext);
    if (!ctx) {
        throw new Error('usePayloadBuilder must be used within a PayloadProvider');
    }
    return ctx;
}