import { OpenAI } from 'openai'
import { OpenAPI } from '../../gen/core/OpenAPI'
import { ResourceService, Script, WorkspaceService } from '../../gen'

import { existsOpenaiKeyStore, workspaceStore } from '$lib/stores'
import { formatResourceTypes } from './utils'
import { scriptLangToEditorLang } from '$lib/scripts'

const WARNING_MSG =
	'WARNING: this code was generated by OpenAI, it might not be accurate. Use at your own risk.'

const COMMENT_TYPES = {
	[Script.language.PYTHON3]: '#',
	[Script.language.DENO]: '//',
	[Script.language.GO]: '//',
	[Script.language.BASH]: '#',
	[Script.language.POSTGRESQL]: '--',
	[Script.language.MYSQL]: '--',
	[Script.language.NATIVETS]: '//',
	[Script.language.BUN]: '//'
}

const PROMPTS = {
	[Script.language
		.PYTHON3]: `Write a function in python called "main". The function should {description}. Specify the parameter types. Do not call the main function.
You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.`,
	[Script.language
		.DENO]: `Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Deno environment. You can import deno libraries or you can also import npm libraries like that: "import ... from "npm:{package}";". Export the "main" function like this: "export function main(...)". Do not call the main function.
You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.`,
	[Script.language.GO]:
		'Write a function in go called "main". The function should {description}. Import the packages you need. The return type of the function has to be ({return_type}, error). The file package has to be "inner".',
	[Script.language.BASH]:
		'Write bash code that should {description}. Do not include "#!/bin/bash". Arguments are always string and can only be obtained with "var1="$1"", "var2="$2"", etc... You do not need to check if the arguments are present.',
	[Script.language.POSTGRESQL]:
		'Write SQL code for a PostgreSQL that should {description}. Arguments can be obtained directly in the statement with $1::{type}, $2::{type}, etc... Name the parameters by adding comments before the command like that: -- $1 name1 (one per row, do not include the type)',
	[Script.language.MYSQL]:
		'Write SQL code for MySQL that should {description}. Arguments can be obtained directly in the statement with ?. Name the parameters by adding comments before the command like that: -- ? name1 ({type}) (one per row)',
	[Script.language
		.NATIVETS]: `Write a function in typescript called "main". The function should {description}. Specify the parameter types. You should use fetch and are not allowed to import any libraries. Export the "main" function like this: "export function main(...)". Do not call the main function.
You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.`,
	[Script.language
		.BUN]: `Write a function in typescript called "main".  The function should {description}. Specify the parameter types. You are in a Node.js environment. You can import npm libraries. Export the "main" function like this: "export function main(...)". Do not call the main function.
You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.`
}

function scriptLangToEnvironment(lang: Script.language) {
	if (lang === Script.language.DENO) {
		return 'typescript in a deno running environment'
	} else if (lang === Script.language.BUN) {
		return 'typescript in a node.js running environment'
	} else if (lang === Script.language.NATIVETS) {
		return 'typescript where you should use fetch and are not allowed to import any libraries'
	} else {
		return lang
	}
}

const SYSTEM_PROMPT =
	'You write code as queried by the user. Only output code. Wrap the code like that: ```language\n{code}\n```. Put any explanation directly in the code as comments.'

const EDIT_SYSTEM_PROMPT =
	'You edit the code as queried by the user. Only output code. Wrap the code like that: ```language\n{code}\n```. Put any explanation directly in the code as comments.'

const EDIT_PROMPT =
	"Here's my environement: {environment}\nHere's my code: ```{lang}\n{code}\n```\nMy instructions: {description}"

const FIX_PROMPT =
	"Here's my environement: {environment}\nHere's my code: ```{lang}\n{code}\n```\nI get the following error: {error}\nFix it for me."

const FIX_SYSTEM_PROMPT =
	'You fix the code shared by the user. Only output code. Wrap the code like that: ```language\n{code}\n```. Put explanations directly in the code as comments.'

export const SUPPORTED_LANGUAGES = new Set(Object.keys(PROMPTS))

let workspace: string | undefined = undefined

workspaceStore.subscribe(async (value) => {
	workspace = value
	if (workspace) {
		try {
			existsOpenaiKeyStore.set(await WorkspaceService.existsOpenaiKey({ workspace }))
		} catch (err) {
			existsOpenaiKeyStore.set(false)
			console.error('Could not get if openai key exists')
		}
	}
})

interface BaseOptions {
	language: Script.language
}

interface ScriptGenerationOptions extends BaseOptions {
	language: Script.language
	description: string
}

interface EditScriptOptions extends ScriptGenerationOptions {
	selectedCode: string
}

interface FixScriptOpions extends BaseOptions {
	code: string
	error: string
}

export async function generateScript(scriptOptions: ScriptGenerationOptions): Promise<string> {
	if (!workspace) {
		throw new Error('No workspace selected')
	}

	const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
	const openai = new OpenAI({
		baseURL,
		apiKey: 'fakekey',
		defaultHeaders: {
			Authorization: ''
		}
	})

	let prompt = PROMPTS[scriptOptions.language].replace('{description}', scriptOptions.description)

	if (
		[Script.language.DENO, Script.language.BUN, Script.language.NATIVETS].includes(
			scriptOptions.language
		)
	) {
		const resourceTypes = await ResourceService.listResourceType({ workspace })
		const resourceTypesText = formatResourceTypes(resourceTypes, 'typescript')

		prompt = prompt.replace('{resourceTypes}', resourceTypesText)
	} else if (scriptOptions.language == Script.language.PYTHON3) {
		const resourceTypes = await ResourceService.listResourceType({ workspace })
		const resourceTypesText = formatResourceTypes(resourceTypes, 'python3')
		prompt = prompt.replace('{resourceTypes}', resourceTypesText)
	}

	const completion = await openai.chat.completions.create({
		model: 'gpt-4',
		max_tokens: 2048,
		messages: [
			{
				role: 'system',
				content: SYSTEM_PROMPT
			},
			{
				role: 'user',
				content: prompt
			}
		]
	})

	let result = completion.choices[0]?.message?.content

	if (!result) {
		throw new Error('No result from OpenAI')
	}

	const match = result.match(/```[a-zA-z]+\n([\s\S]*?)\n```/)

	if (!match || match.length < 2) {
		throw new Error('No code block found')
	}

	result = match[1]

	if (scriptOptions.language == Script.language.GO) {
		const warning = COMMENT_TYPES[scriptOptions.language] + ' ' + WARNING_MSG + '\n'

		return result.trim().replace('package inner\n', 'package inner\n' + warning)
	} else if (scriptOptions.language == Script.language.BASH) {
		return (
			'# shellcheck shell=bash\n' +
			COMMENT_TYPES[scriptOptions.language] +
			' ' +
			WARNING_MSG +
			'\n\n' +
			result.trim()
		)
	} else {
		return COMMENT_TYPES[scriptOptions.language] + ' ' + WARNING_MSG + '\n\n' + result.trim()
	}
}

export async function editScript(scriptOptions: EditScriptOptions): Promise<string> {
	if (!workspace) {
		throw new Error('No workspace selected')
	}

	const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
	const openai = new OpenAI({
		baseURL,
		apiKey: 'fakekey',
		defaultHeaders: {
			Authorization: ''
		}
	})

	let prompt = EDIT_PROMPT.replace('{lang}', scriptLangToEditorLang(scriptOptions.language))
		.replace('{code}', scriptOptions.selectedCode)
		.replace('{description}', scriptOptions.description)
		.replace('{environment}', scriptLangToEnvironment(scriptOptions.language))

	const completion = await openai.chat.completions.create({
		model: 'gpt-4',
		max_tokens: 512,
		messages: [
			{
				role: 'system',
				content: EDIT_SYSTEM_PROMPT
			},
			{
				role: 'user',
				content: prompt
			}
		],
		temperature: 0.5
	})

	let result = completion.choices[0]?.message?.content

	if (!result) {
		throw new Error('No result from OpenAI')
	}

	const match = result.match(/```[a-zA-z]+\n([\s\S]*?)\n```/)

	if (!match || match.length < 2) {
		throw new Error('No code block found')
	}

	result = match[1]

	return result
}

export async function fixScript(scriptOptions: FixScriptOpions) {
	if (!workspace) {
		throw new Error('No workspace selected')
	}

	const baseURL = `${location.origin}${OpenAPI.BASE}/w/${workspace}/openai/proxy`
	const openai = new OpenAI({
		baseURL,
		apiKey: 'fakekey',
		defaultHeaders: {
			Authorization: ''
		}
	})

	let prompt = FIX_PROMPT.replace('{lang}', scriptLangToEditorLang(scriptOptions.language))
		.replace('{code}', scriptOptions.code)
		.replace('{error}', scriptOptions.error)
		.replace('{environment}', scriptLangToEnvironment(scriptOptions.language))

	const completion = await openai.chat.completions.create({
		model: 'gpt-4',
		max_tokens: 512,
		messages: [
			{
				role: 'system',
				content: FIX_SYSTEM_PROMPT
			},
			{
				role: 'user',
				content: prompt
			}
		]
	})

	let result = completion.choices[0]?.message?.content

	if (!result) {
		throw new Error('No result from OpenAI')
	}

	const match = result.match(/```[a-zA-z]+\n([\s\S]*?)\n```/)

	if (!match || match.length < 2) {
		throw new Error('No code block found')
	}

	result = match[1]

	return result
}