from fastapi import FastAPI, BackgroundTasks
from fastapi.responses import JSONResponse, HTMLResponse
from pydantic import BaseModel
import uvicorn
from uvicorn.config import logger
import os
import argparse
from tasks.training import TrainingTask
from config import enable_test_mode

app = FastAPI()

@app.post("/train/start", response_class=JSONResponse)
async def start_model_training(background_tasks: BackgroundTasks):
	"""
	Endpoint on which a request can be sent to start model re-training, 
	if there's no training task currently running.
	The task will be carried out in background and its status can be
	polled via /train/get_state.

	Args:
		background_tasks (BackgroundTasks): BG Tasks scheduler provided by FastAPI

	Returns:
		dict: A dictionary containing a message of the outcome for the request.
	"""
	if not TrainingTask.has_instance():
		background_tasks.add_task(TrainingTask.get_instance())

		return {
			"message": "Model training was scheduled and will begin shortly.",
		}
	
	return {
		"message": "A training instance is already running.",
	}

@app.post("/train/get_state", response_class=JSONResponse)
async def poll_model_training_state():
	"""
	Checks if there is currently a training task ongoing.
	If so, returns whether it's done and/or if an error occurred.
	Otherwise if no instance is running, returns only a message.

	Returns:
		dict: Dictionary containing either done/error or message.
	"""
	if TrainingTask.has_instance():
		train_instance : TrainingTask = TrainingTask.get_instance()
		is_done = train_instance.is_done()
		has_error = train_instance.has_error()

		if is_done:
			TrainingTask.clear_instance()

		return {
			"done": is_done,
			"error": has_error,
		}

	return {
		"message": "No training instance running!",
	}

class InferenceRequest(BaseModel):
	"""
	Provides a model/schema for the accepted request body for incoming
	inference requests.
	"""
	messages: list[str]

@app.post("/inference", response_class=JSONResponse)
async def inference(data: InferenceRequest):
	"""
	Endpoint on which you can send a list of messages that need to be
	classified. A list of predictions will be returned in response,
	containing for each message all of the probabilities for each label.

	Args:
		data (InferenceRequest): Structure containing a list of 
  								 messages that shall be evaluated

	Returns:
		json: A json list containing the sentiment analysis for each message.
			  Each element consists of a dictionary with the following keys:
			  positive, neutral, negative
	"""
	
	from tasks.inference import infer_task
	return infer_task.predict(data.messages)

@app.get("/", response_class=HTMLResponse)
async def root():
	"""
	The root endpoint for our hosted application. Only shows a message
	showing that it's up and running.

	Returns:
		str: A html response containing a hello world-like string
	"""
	return "Hi there! It's a nice blank page, isn't it?"

if __name__ == "__main__":
	"""
	Entrypoint for the application executed via command-line.
	It accepts an optional argument "--test" to enable the test mode.
	"""
	parser = argparse.ArgumentParser()
	parser.add_argument("test", nargs="?", default="no")
	args = parser.parse_args()

	if args.test == "yes":
		enable_test_mode()

	config = uvicorn.Config("main:app", host="0.0.0.0", port=int(os.environ["APP_LISTEN_PORT"]), log_level="debug")
	server = uvicorn.Server(config)
	server.run()