from typing import Any, Dict, List
 
import torch
import json
import transformers
import bitsandbytes
from transformers import AutoModelForCausalLM, AutoTokenizer
 
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
 
class EndpointHandler:
    def __init__(self, path=""):
        print("ANSHAY calling EndpointHandler __init__")
        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            return_dict=True,
            device_map="auto",
            load_in_8bit=True,
            torch_dtype=dtype,
            trust_remote_code=True,
        )

        
        print("ANSHAY got model and tokenizer")
        print(self.model)
        print(self.tokenizer)
 
        generation_config = self.model.generation_config
        generation_config.max_new_tokens = 3
        generation_config.temperature = 0.1
        generation_config.pad_token_id = self.tokenizer.eos_token_id
        generation_config.eos_token_id = self.tokenizer.eos_token_id
        self.generation_config = generation_config
 
        self.pipeline = transformers.pipeline(
            "text-generation", model=self.model, tokenizer=self.tokenizer
        )

        print("ANSHAY got pipeline")
        print(self.pipeline)
 
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        print("ANSHAY in __call__ of handler")
        print("ANSHAY printing data:")
        print(data)
        
        input_string = data.pop("inputs");
        print("ANSHAY printing input_string:")
        print(input_string)
        
        input_dict = json.loads(input_string)
        print("ANSHAY printing input_dict:")
        print(input_dict)
        
        gender = input_dict["gender"]
        age = input_dict["age"]
        occupation = input_dict["occupation"]
        sleep_hours = input_dict["sleep_hours"]
        physical_activity = input_dict["physical_activity"]
        caffiene_intake = input_dict["caffiene_intake"]
        alcohol_consumption = input_dict["alcohol_consumption"]
        smoking = input_dict["smoking"]
        family_anxiety = input_dict["family_anxiety"]
        stress_level = input_dict["stress_level"]
        heart_rate = input_dict["heart_rate"]
        breathing_rate = input_dict["breathing_rate"]
        sweating_level = input_dict["sweating_level"]
        dizziness = input_dict["dizziness"]
        medication = input_dict["medication"]
        therapy_sessions = input_dict["therapy_sessions"]
        recent_life_event = input_dict["recent_life_event"]
        diet_quality = input_dict["diet_quality"]

        chat_formatted = [{
            "role": "system",
            "content": "You are an anxiety calculating bot. You help people rate their anxiety on a scall of 1-10, after understanding their current lifestyle.",
        },
        {
            "role": "user",
            "content": f"Here are my details:\nGender:{gender}\nAge:{age}\nOccupation:{occupation}\nSleep Hours:{sleep_hours}\nPhysical Activity (hrs/week):{physical_activity}\nCaffeine Intake (mg/day):{caffiene_intake}\nAlcohol Consumption (drinks/week):{alcohol_consumption}\nSmoking:{smoking}\nFamily History of Anxiety:{family_anxiety}\nStress Level (1-10):{stress_level}\nHeart Rate (bpm):{heart_rate}\nBreathing Rate (breaths/min):{breathing_rate}\nSweating Level (1-5):{sweating_level}\nDizziness:{dizziness}\nMedication:{medication}\nTherapy Sessions (per month):{therapy_sessions}\nRecent Major Life Event:{recent_life_event}\nDiet Quality (1-10):{diet_quality}"
        },
        {
            "role": "assistant",
            "content": f"Anxiety Rating:"
        }]

        print("ANSHAY printing final chat formatted:")
        print(chat_formatted)

        chat_template_applied = self.tokenizer.apply_chat_template(chat_formatted, tokenize=False, continue_final_message=True)

        print("ANSHAY printing final chat_template_applied:")
        print(chat_template_applied)
        
        result = self.pipeline(chat_template_applied, generation_config=self.generation_config)
        print("ANSHAY printing final result:")
        print(result)
        
        generated_text = result[0]['generated_text']
        print("ANSHAY printing generated_text:")
        print(generated_text)
        
        anxiety_rating = generated_text[-3:]
        print("ANSHAY printing anxiety_rating:")
        print(anxiety_rating)
        
        return anxiety_rating