import os
import requests
from enum import Enum
from dotenv import load_dotenv
from prompts import Prompts
import json

# Load environment variables from .env file
load_dotenv()

# Access environment variables
API_KEY = os.environ['openAI_key']

class Stage(Enum):
    INITIAL_ASSESSMENT = "Initial Assessment"
    FOLLOWUP_ASSESSMENT = "Follow Up Assessment"
    EVALUATION = "Evaluation"
    DETAILED_EVALUATION = "Detailed Evaluation"
    PROGRESS_NOTE = "Progress Note"
    RECOMMENDATION="Recommendations"
    DISCHARGE_SUMMARY = "Discharge Summary"
    SHORT_TERM_GOALS = "Short Term Goals"
    LONG_TERM_GOALS = "Long Term Goals"


class GPT_Service:
    def __init__(self):
        self.api_key = API_KEY
        self.base_url = 'https://api.openai.com/v1/chat/completions'

    def request_gpt_for_response(self, system_prompt, userPrompt):
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        data = {
            'model': 'gpt-3.5-turbo-1106',
            'temperature': 0.3,
            'messages': [
                {
                    'role': 'system',
                    'content': system_prompt
                },
                {
                    'role': 'user',
                    'content': userPrompt
                }
            ]
        }
        try:
            response = requests.post(
                self.base_url,
                json=data,
                headers=headers,
            )
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Error: {e}")
    
    def formulateUserPrompt(self, assessmentPhase, keywords, medicalHistory):
        patientData= f"""
        Keywords: {keywords}.
        {f"Patient has the following medical history: {medicalHistory}" if medicalHistory is not None and medicalHistory != ""
        else ""}
        """
        print(patientData)
        return patientData
    
    def processOutputResults(self, jsonOutput):
        preprocessed_json_string = jsonOutput.replace('\n', '')
        data = json.loads(preprocessed_json_string)
        dataLength=len(data.keys())
        result_list = []
        
        for key, value in data.items():
            result_list.append(key)
            for variant, sentence in value.items():
                result_list.append(sentence)
                
        if dataLength < 5:
            additional_keys_needed = 5 - dataLength
            for key in range(additional_keys_needed):
                result_list.append(f"Keyword {dataLength+key+1}")  # Empty key
                for _ in range(3):
                    result_list.append("")  # Empty sentence
            return result_list

        return result_list

    def getInitialAssessment(self, keywords,medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.INITIAL_ASSESSMENT.value,keywords,medicalHistory)
        response= self.request_gpt_for_response(Prompts.INITIAL_ASSESSMENT_PROMPT.value,userPrompt)
        return self.processOutputResults(response)

    def getFollowUpAssessment(self, keywords, medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.FOLLOWUP_ASSESSMENT.value,keywords,medicalHistory)
        response= self.request_gpt_for_response(Prompts.FOLLOWUP_ASSESSMENT_PROMPT.value,userPrompt)
        return self.processOutputResults(response)

    def getEvaluation(self, keywords, medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.EVALUATION.value,keywords,medicalHistory)
        response= self.request_gpt_for_response(Prompts.EVALUATION_PROMPT.value,userPrompt)
        return self.processOutputResults(response)
    
    def getRecommendation(self, keywords, medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.RECOMMENDATION.value,keywords,medicalHistory)
        response= self.request_gpt_for_response(Prompts.RECOMMENDATION_PROMPT.value,userPrompt)
        return self.processOutputResults(response)

    def getDetailedEvaluation(self, keywords,medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.DETAILED_EVALUATION.value,keywords,medicalHistory)
        response= self.request_gpt_for_response(Prompts.DETAILED_EVALUATION_PROMPT.value,userPrompt)
        return self.processOutputResults(response)

    def getProgressNote(self, keywords,medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.PROGRESS_NOTE.value,keywords,medicalHistory)
        response = self.request_gpt_for_response(Prompts.PROGRESS_NOTE_PROMPT.value,userPrompt)
        return self.processOutputResults(response)

    def getDischargeSummary(self, keywords,medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.DISCHARGE_SUMMARY.value,keywords,medicalHistory)
        response = self.request_gpt_for_response(Prompts.DISCHARGE_SUMMARY_PROMPT.value,userPrompt)
        return self.processOutputResults(response)
    
    def getShortTermGoals(self, keywords,medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.SHORT_TERM_GOALS.value,keywords,medicalHistory)
        response =  self.request_gpt_for_response(Prompts.SHORT_TERM_GOALS_PROMPT.value,userPrompt)
        return self.processOutputResults(response)
    
    def getLongTermGoals(self, keywords,medicalHistory):
        userPrompt= self.formulateUserPrompt(Stage.LONG_TERM_GOALS.value,keywords,medicalHistory)
        response = self.request_gpt_for_response(Prompts.LONG_TERM_GOALS_PROMPT.value,userPrompt)
        return self.processOutputResults(response)
    
    def getRephrasedSentences(self, sentences):
        response = self.request_gpt_for_response(Prompts.REPHRASE_PROMPT.value,sentences)
        return response