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):
        print(f"System Prompt: {system_prompt}\n user Prompt: {userPrompt}\n")
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        data = {
            'model': 'gpt-4-turbo',
            '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 ""}
        """
        return patientData
    
    def formulateUserPromptForDetailAssessment(self, rawPatientInformation, medicalHistory):
        patientData= f"""
        Sentences: {rawPatientInformation}.
        {f"Patient has the following medical history: {medicalHistory}" if medicalHistory is not None and medicalHistory != ""
        else ""}
        """
        return patientData
    
    def processOutputResults(self, jsonOutput):
        preprocessed_json_string = jsonOutput.replace('\n', '')
        data = json.loads(preprocessed_json_string)
        result_list = []
        for key, value in data.items():
            for variant, sentence in value.items():
                result_list.append(sentence)

        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 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 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)
    
    # Detail Patient Evaluation Prompt

    def getDetailInitialAssessment(self,medicalHistory,sentences):
        rawPatientAssessment = " ".join(sentences)
        userPrompt= self.formulateUserPromptForDetailAssessment(rawPatientAssessment,medicalHistory)
        response= self.request_gpt_for_response(Prompts.DETAIL_PATIENT_ASSESSMENT_INITIAL_ASSESSMENT_GOALS.value,userPrompt)
        return response
    
    def getDetailShortTermGoals(self,medicalHistory,sentences):
        rawPatientAssessment = " ".join(sentences)
        userPrompt= self.formulateUserPromptForDetailAssessment(rawPatientAssessment,medicalHistory)
        response= self.request_gpt_for_response(Prompts.DETAIL_PATIENT_ASSESSMENT_SHORT_TERM_GOALS.value,userPrompt)
        return response

    def getDetailLongTermGoals(self,medicalHistory,sentences):
        rawPatientAssessment = " ".join(sentences)
        userPrompt= self.formulateUserPromptForDetailAssessment(rawPatientAssessment,medicalHistory)
        response= self.request_gpt_for_response(Prompts.DETAIL_PATIENT_ASSESSMENT_LONG_TERM_GOALS.value,userPrompt)
        return response
    
    def getDetailRecommendation(self, medicalHistory,sentences):
        rawPatientAssessment = " ".join(sentences)
        userPrompt= self.formulateUserPromptForDetailAssessment(rawPatientAssessment,medicalHistory)
        response= self.request_gpt_for_response(Prompts.DETAIL_PATIENT_ASSESSMENT_RECOMMENDATION_GOALS.value,userPrompt)
        return response
    

    # Extra helper functions
    
    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 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)