import openai
import pyjson5


class Env:
	def __init__(self):
		self.model = "gpt-4o"
		self.output_tokens_min = 250
		self.output_tokens_max = 500
		self.temp: float = None
		self.seed: int = None
		self.max_tokens = 8000

class SourceFile:
	def __init__(self):
		self.name: str = None
		self.category: str = None
		self.contents: str = None

class LLMAPIException (Exception): ...

def __do_prompt (lang: str, extra_prompt: str, env: Env):
	prompt = '''
Write me a program source code written in %s. Write anything you'd like.
Use more than %d words, but no more than %d words in the code.''' % (
			   lang, env.output_tokens_min, env.output_tokens_max)
	if extra_prompt:
		prompt = ' ' + extra_prompt

	messages = [
		{
			"role": "system",
			"content": "You're a helpful assistant that writes "
				+ "computer program of any kind"
		},
		{ "role": "user", "content": prompt }
	]
	functions = [
		{
			"name": "get_result",
			"description": "Output the source code",
			"parameters": {
				"type": "object",
				"properties": {
					"cat": {
						"type": "string",
						"description": "one word describing program functions"
					},
					"filename": {
						"type": "string",
						"description": "The source code file name" # TODO: strip unallowed characters in Unix fs
					},
					"code": {
						"type": "string",
						"description": "The source code"
					}
				}
			}
		}
	]

	return openai.chat.completions.create(
		model = env.model,
		messages = messages,
		functions = functions,
		function_call = "auto",
		temperature = env.temp,
		seed = env.seed,
		max_tokens = env.max_tokens
	)

def gen_rand_srccode (lang: str, extra_prompt: str, env: Env) -> SourceFile | str:
	rsp = __do_prompt(lang, extra_prompt, env)
	match rsp.choices[0].finish_reason:
		case 'stop':
			return rsp.choices[0].message.content
		case 'function_call':
			choice = pyjson5.loads(rsp.choices[0].message.function_call.arguments)

			ret = SourceFile()
			ret.name = choice["filename"]
			ret.category = choice.get("cat")
			ret.contents = choice["code"]

			return ret
		case _: raise LLMAPIException(rsp.choices[0].finish_reason)