import { assert } from '@takeshape/util';
import set from 'lodash/fp/set.js';
import get from 'lodash/get.js';
import uniq from 'lodash/uniq.js';
import upperFirst from 'lodash/upperFirst.js';
import * as abortAgentMessageResponse from "./resolvers/ai/abort-agent-message-response.js";
import * as agentChatPayload from "./resolvers/ai/agent-chat-payload.js";
import * as agentSessionPayload from "./resolvers/ai/agent-session-payload.js";
import * as getAgentMessageArgs from "./resolvers/ai/get-agent-message-args.js";
import * as getAgentMessageResponse from "./resolvers/ai/get-agent-message-response.js";
import * as inspectAgentArgs from "./resolvers/ai/inspect-agent-args.js";
import * as inspectAgentResponse from "./resolvers/ai/inspect-agent-response.js";
import * as sendAgentFeedbackArgs from "./resolvers/ai/send-agent-feedback-args.js";
import * as sendAgentMessageArgs from "./resolvers/ai/send-agent-message-args.js";
import * as sendAgentMessageResponse from "./resolvers/ai/send-agent-message-response.js";
import { JSON_SHAPE_NAME } from "./resolvers/takeshape/builtins/constants.js";
/**
 * Special transition destination ID for a newly created, never started session.
 */
export const AGENT_SESSION_CREATED = 'SESSION_CREATED';
/**
 * Special transition destination ID for a started session.
 */
export const START_AGENT_EXECUTION = 'START';
/**
 * Special transition destination ID for agent transitions that end the agent execution.
 */
export const END_AGENT_EXECUTION = 'END';
/**
 * Special history strategy ID that indicates the history should be cleared.
 */
export const CLEAR_HISTORY_STRATEGY_ID = 'clearHistory';
export const EXECUTION_TYPE_CHAT = 'chat';
export const EXECUTION_TYPE_GENERATE = 'generate';
export const EXECUTION_TYPE_GRAPHQL = 'graphql';
function canEndOnState(state, includeSuspend = false) {
    const { transitions } = state;
    if (!transitions?.length) {
        return true;
    }
    for (const step of transitions) {
        if (includeSuspend && step.type === 'suspend') {
            return true;
        }
        if (step.destination === END_AGENT_EXECUTION) {
            return true;
        }
        if (step.condition === undefined) {
            return false;
        }
    }
    return true;
}
export function getAgentEndStates(agent, includeSuspend = false) {
    const endStates = new Set();
    for (const [stateId, state] of Object.entries(agent.states)) {
        if (canEndOnState(state, includeSuspend)) {
            endStates.add(stateId);
        }
    }
    return endStates;
}
// Agent Scoped Names
export const getCreateAgentSessionMutationName = (agentName) => {
    return `create${upperFirst(agentName)}Session`;
};
export const getInspectAgentQueryName = (agentName) => {
    return `inspect${upperFirst(agentName)}`;
};
export const getSendAgentFeedbackMutationName = (agentName) => {
    return `send${upperFirst(agentName)}Feedback`;
};
export const getGetAgentMessageQueryName = (agentName) => {
    return `get${upperFirst(agentName)}Message`;
};
export const getAbortAgentMessageMutationName = (agentName) => {
    return `abort${upperFirst(agentName)}Message`;
};
// Agent Input Scopes Names
export const getSendAgentMessageMutationName = (inputName) => {
    return `send${upperFirst(inputName)}Message`;
};
function createAssertValidName(key, result) {
    const singular = key === 'shapes' ? 'shape' : key === 'mutations' ? 'mutation' : 'query';
    return (name) => {
        assert(result[key][name] === undefined, `Schema already has a ${singular} with the name "${name}"`);
    };
}
export function getAgentFeatures(agents) {
    const result = {
        queries: {},
        mutations: {},
        shapes: {}
    };
    const assertValidShapeName = createAssertValidName('shapes', result);
    const assertValidQueryName = createAssertValidName('queries', result);
    const assertValidMutationName = createAssertValidName('mutations', result);
    for (const [agentName, agent] of Object.entries(agents)) {
        // Get valid return types based on states that could possibly be the end state
        // TODO In the future the agent will have a return type defined and we will instead validate
        // that all the return states will return the correct shape.
        const returnStates = [...getAgentEndStates(agent, true)];
        const returnTypes = uniq(returnStates.map((stateId) => {
            const { execution } = agent.states[stateId];
            if (execution.type === EXECUTION_TYPE_CHAT) {
                return agentChatPayload.shapeName;
            }
            if (execution.type === EXECUTION_TYPE_GENERATE) {
                return execution.outputShape ? execution.outputShape : 'string';
            }
            return JSON_SHAPE_NAME;
        }));
        const agentOutput = returnTypes.length === 0 ? 'string' : returnTypes.length === 1 ? returnTypes[0] : JSON_SHAPE_NAME;
        const agentOutputSchema = agentOutput === 'string' ? { type: 'string' } : { '@ref': agentOutput };
        const agentSessionShapeName = agentSessionPayload.shapeName;
        const agentMessageShapeName = getAgentMessageResponse.createUserShapeName(agentName);
        assertValidShapeName(agentMessageShapeName);
        result.shapes[agentMessageShapeName] = {
            id: agentMessageShapeName,
            name: agentMessageShapeName,
            title: agentMessageShapeName,
            schema: getAgentMessageResponse.createUserShapeSchema({
                properties: {
                    output: agentOutputSchema
                }
            })
        };
        const sendAgentMessageResponseShapeName = sendAgentMessageResponse.createUserShapeName(agentName);
        assertValidShapeName(sendAgentMessageResponseShapeName);
        result.shapes[sendAgentMessageResponseShapeName] = {
            id: sendAgentMessageResponseShapeName,
            name: sendAgentMessageResponseShapeName,
            title: sendAgentMessageResponseShapeName,
            schema: sendAgentMessageResponse.createUserShapeSchema({
                properties: {
                    output: agentOutputSchema
                }
            })
        };
        const inspectAgentResponseShapeName = inspectAgentResponse.createUserShapeName(agentName);
        assertValidShapeName(inspectAgentResponseShapeName);
        result.shapes[inspectAgentResponseShapeName] = {
            id: inspectAgentResponseShapeName,
            name: inspectAgentResponseShapeName,
            title: inspectAgentResponseShapeName,
            schema: inspectAgentResponse.createUserShapeSchema()
        };
        const getAgentMessageQueryName = getGetAgentMessageQueryName(agentName);
        assertValidQueryName(getAgentMessageQueryName);
        result.queries[getAgentMessageQueryName] = {
            description: `Get a message for the ${agentName} agent`,
            shape: agentMessageShapeName,
            resolver: {
                name: 'ai:getAgentMessage',
                agentName
            },
            args: {
                ...getAgentMessageArgs.schema
            }
        };
        const abortAgentMessageMutationName = getAbortAgentMessageMutationName(agentName);
        assertValidMutationName(abortAgentMessageMutationName);
        result.mutations[abortAgentMessageMutationName] = {
            description: `Abort a message for the ${agentName} agent`,
            shape: abortAgentMessageResponse.shapeName,
            resolver: {
                name: 'ai:abortAgentMessage',
                agentName
            },
            args: {
                ...getAgentMessageArgs.schema
            }
        };
        const createAgentSessionMutationName = getCreateAgentSessionMutationName(agentName);
        assertValidMutationName(createAgentSessionMutationName);
        result.mutations[createAgentSessionMutationName] = {
            description: `Create a session for the ${agentName} agent`,
            shape: agentSessionShapeName,
            resolver: {
                name: 'ai:createAgentSession',
                agentName
            }
        };
        const inspectAgentQueryName = getInspectAgentQueryName(agentName);
        assertValidQueryName(inspectAgentQueryName);
        result.queries[inspectAgentQueryName] = {
            description: `Inspect the ${agentName} agent`,
            args: inspectAgentArgs.schema,
            shape: inspectAgentResponseShapeName,
            resolver: {
                name: 'ai:inspectAgent',
                agentName
            }
        };
        result.mutations[getSendAgentFeedbackMutationName(agentName)] ||= {
            description: `Send feedback for the ${agentName} agent`,
            args: sendAgentFeedbackArgs.shapeName,
            shape: 'boolean',
            resolver: {
                name: 'ai:sendAgentFeedback',
                agentName
            }
        };
        // This mutation can vary depending on the inputs defined, everything else is stable for the whole agent (for now)
        for (const inputConfig of agent.api.inputs) {
            if (inputConfig.type === 'mutation') {
                const inputName = inputConfig.name;
                const sendAgentMessageMutationName = getSendAgentMessageMutationName(inputName);
                assertValidMutationName(sendAgentMessageMutationName);
                result.mutations[sendAgentMessageMutationName] = {
                    description: agent.description,
                    args: {
                        type: 'object',
                        extends: [
                            {
                                ...sendAgentMessageArgs.shapeSchema,
                                required: agent.protection?.enabled
                                    ? [...sendAgentMessageArgs.shapeSchema.required, 'token']
                                    : sendAgentMessageArgs.shapeSchema.required
                            },
                            {
                                '@ref': inputConfig.args
                            }
                        ]
                    },
                    shape: sendAgentMessageResponseShapeName,
                    resolver: {
                        name: 'ai:sendAgentMessage',
                        agentName,
                        inputName
                    }
                };
            }
        }
    }
    return result;
}
export function isValidAgentMutation(projectSchema, name) {
    for (const agent of Object.values(getAgentMap(projectSchema) ?? {})) {
        for (const input of agent.api.inputs) {
            if (input.type === 'mutation' && input.name === name) {
                return true;
            }
        }
    }
    return false;
}
export function addAiQueries(projectSchema) {
    const agents = getAgentMap(projectSchema);
    if (!agents) {
        return projectSchema;
    }
    const aiFeatures = getAgentFeatures(agents);
    return {
        ...projectSchema,
        queries: {
            ...projectSchema.queries,
            ...aiFeatures.queries
        },
        mutations: {
            ...projectSchema.mutations,
            ...aiFeatures.mutations
        },
        shapes: {
            ...projectSchema.shapes,
            ...aiFeatures.shapes
        }
    };
}
/**
 * Where guards live in the schema.
 */
export const GUARDS_SCHEMA_PATH = ['guards'];
export const AGENTS_SCHEMA_PATH = ['agents'];
/**
 * Check if a guard is enabled.
 */
export const isGuardEnabled = (guard) => {
    return Boolean(guard && guard.enabled !== false);
};
export function getGuardMap(schema) {
    return schema.guards;
}
export function getGuards(schema) {
    const guards = getGuardMap(schema);
    return guards ? Object.entries(guards).map(([id, guard]) => ({ id, ...guard })) : [];
}
export function setGuardMap(schema, guards) {
    return set(GUARDS_SCHEMA_PATH, guards, schema);
}
export function setGuard(schema, id, guard) {
    return set([...GUARDS_SCHEMA_PATH, id], guard, schema);
}
/**
 * Get a guard configuration from the project schema.
 */
export function getGuardConfig(schema, guardId) {
    return get(schema, [...GUARDS_SCHEMA_PATH, guardId]);
}
/**
 * Get a guard configuration from the project schema, but only if it is enabled.
 */
export function getEnabledGuardConfig(schema, guardId) {
    const config = getGuardConfig(schema, guardId);
    if (config && isGuardEnabled(config)) {
        return config;
    }
}
export function getAgentMap(projectSchema) {
    return projectSchema.agents;
}
export function getAgents(projectSchema) {
    const agentMap = getAgentMap(projectSchema);
    return agentMap && Object.entries(agentMap).map(([name, agent]) => ({ name, ...agent }));
}
export function setAgentMap(projectSchema, agents) {
    return set(AGENTS_SCHEMA_PATH, agents, projectSchema);
}
export function setAgent(projectSchema, name, agent) {
    return set([...AGENTS_SCHEMA_PATH, name], agent, projectSchema);
}
export function getAgent(projectSchema, name) {
    return getAgentMap(projectSchema)?.[name];
}
export function ensureAgent(projectSchema, name) {
    const agent = getAgent(projectSchema, name);
    if (!agent) {
        throw new Error(`Agent "${name}" not found`);
    }
    return agent;
}
export function isAgentQuery(query, agentName) {
    const resolver = query.resolver;
    if ('agentName' in resolver) {
        return resolver.agentName === agentName;
    }
    return false;
}
