#!/usr/bin/env python
# coding: utf-8
from os import listdir
from os.path import isdir
from fastapi import FastAPI, HTTPException, Request, responses
from fastapi.middleware.cors import CORSMiddleware
from llama_cpp import Llama
from fastapi import Body

from pydantic import BaseModel
from enum import Enum
from typing import Optional

print("Loading model...")
llm = Llama(
      model_path="/models/final-gemma2b_SA-Q8_0.gguf",
      # n_gpu_layers=28, # Uncomment to use GPU acceleration
      # seed=1337, # Uncomment to set a specific seed
      # n_ctx=2048, # Uncomment to increase the context window
)

def ask(question, max_new_tokens=200):
  output = llm(
    question, # Prompt
    max_tokens=max_new_tokens, # Generate up to 32 tokens, set to None to generate up to the end of the context window
    stop=["\n"], # Stop generating just before the model would generate a new question
    echo=False, # Echo the prompt back in the output
    temperature=0.0,
  )
  return output

def check_sentiment(text):
  result = ask(f'Analyze the sentiment of the tweet enclosed in square brackets, determine if it is positive or negative, and return the answer as the corresponding sentiment label "positive" or  "negative" [{text}] =', max_new_tokens=3)
  return result['choices'][0]['text'].strip()

def clean_sentiment_response(response_text):
  result = response_text.strip()
  if "positive" in result:
    return "positive"
  elif "negative" in result:
    return "negative"
  else:
    return "unknown"
  
print("Testing model...")
assert "positive" in check_sentiment("ดอกไม้ร้านนี้สวยจัง")
print("Ready.")

app = FastAPI(
    title = "GemmaSA_2b",
    description="A simple sentiment analysis API for the Thai language, powered by a finetuned version of Gemma-2b",
    version="1.0.0",
)

origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

class SA_Result(str, Enum):
  positive = "positive"
  negative = "negative"
  unknown = "unknown"

class SA_Response(BaseModel):
  text: Optional[str] = None
  code: int = 200
  result: SA_Result = None



@app.get('/')
def docs():
  "Redirects the user from the main page to the docs."
  return responses.RedirectResponse('./docs')

@app.get('/add/{a}/{b}')
def add(a: int,b: int):
    return a + b

@app.post('/SA')
def perform_sentiment_analysis(prompt: str = Body(..., embed=True, example="I like eating fried chicken")) -> SA_Response:
  """Performs a sentiment analysis using a finetuned version of Gemma-7b"""
  if prompt:
    try:
      print(f"Checking sentiment for {prompt}")
      result = clean_sentiment_response(check_sentiment(prompt))
      print(f"Result: {result}")
      return SA_Response(result=result, text=prompt)
    except Exception as e:
      return HTTPException(500, SA_Response(code=500, result=str(e), text=prompt))
  else:
    return HTTPException(400, SA_Response(code=400, result="Request argument 'prompt' not provided."))