from decimal import Decimal
from typing import Tuple, Union, Optional

from h_adminsim.registry import STATUS_CODES
from h_adminsim.environment.hospital import HospitalEnvironment
from h_adminsim.utils.common_utils import *



class SanityChecker:
    def __init__(self,
                 start_hour: Optional[float] = None,
                 end_hour: Optional[float] = None,
                 time_unit: Optional[float] = None,):
        
        # Initialization
        self._START_HOUR = start_hour
        self._END_HOUR = end_hour
        self._TIME_UNIT = time_unit


    def intake_check(self,
                     prediction: dict, 
                     gt: dict,
                     conversations: str) -> Tuple[bool, str]:
        """
        Performs a sanity check on the predicted patient information and department against the ground truth.

        Args:
            prediction (dict): The output generated by the model. Expected to contain:
                - 'patient': dict of patient demographic information (e.g., name, birth date, gender)
                - 'department': str representing the predicted department
            gt (dict): The ground truth data. Expected to contain:
                - 'patient': dict of correct patient demographic information
                - 'department': str of the correct department
            conversations (str): The full conversation text between the patient and administration staff.

        Returns:
            Tuple[bool, str]:
                - bool: True if the prediction passes all sanity checks, False otherwise
                - str: Status code indicating the type of check passed or failed
        """
        ############################ Check the prediciton format #############################
        if not isinstance(prediction['patient'], dict):
            return False, STATUS_CODES['format']  # Could not be parsed as a dictionary
        
        ############################ Incomplete simulation case #############################
        if not all(v.lower() in conversations.lower() for k, v in gt['patient'].items()):
            return False, STATUS_CODES['simulation']
        
        ############################ Check with the ground truth #############################
        wrong_department = prediction['department'][0] not in gt['department']
        wrong_info = prediction['patient'] != gt['patient']
        if wrong_department and wrong_info:
            return False, STATUS_CODES['department & patient']
        elif wrong_department:
            return False, STATUS_CODES['department']
        elif wrong_info:
            return False, STATUS_CODES['patient']
        
        return True, STATUS_CODES['correct']
    

    def __check_is_earliest(self, 
                            prediction: dict, 
                            gt_patient_condition: dict, 
                            doctor_information: dict, 
                            environment: HospitalEnvironment) -> bool:
        """
        Check if the predicted schedule is the earliest possible option.

        Args:
            prediction (dict): Predicted schedule information including doctor, start time, end time, and date.
            gt_patient_condition (dict): Ground truth patient conditions used only for sanity checks.
            doctor_information (dict): Dictionary containing doctors' schedules and availability.
            environment (HospitalEnvironment): Environment object containing current time and UTC offset.

        Returns:
            bool: True if the predicted schedule is the earliest available, False otherwise.
        """
        # Init grount thruth values
        department = gt_patient_condition['department']
        preference_type = gt_patient_condition['preference']
        valid_from = gt_patient_condition['valid_from']
        fixed_schedules = environment.get_doctor_schedule(doctor_information=doctor_information, department=department)['doctor']

        # Get predicted results
        pred_doctor_name = list(prediction['schedule'].keys())[0]
        pred_start = prediction['schedule'][pred_doctor_name]['start']
        pred_end = prediction['schedule'][pred_doctor_name]['end']
        pred_date = prediction['schedule'][pred_doctor_name]['date']
        current_time = environment.current_time
        utc_offset = environment._utc_offset

        # Time segments
        prediction_schedule_segments = convert_time_to_segment(self._START_HOUR,
                                                               self._END_HOUR,
                                                               self._TIME_UNIT,
                                                               [pred_start, pred_end])
        
        for k, v in fixed_schedules.items():
            if preference_type == 'doctor' and k != pred_doctor_name:
                continue
            
            min_time_slot_n = int(Decimal(str(v['outpatient_duration'])) / Decimal(str(self._TIME_UNIT)))
            fixed_schedule = v['schedule']
            for date, schedule in fixed_schedule.items():
                # date > pred_date case
                if compare_iso_time(date, pred_date):
                    continue
                
                # valid_from > date case (preference == 'date' case)
                if valid_from and compare_iso_time(valid_from, date):
                    continue

                fixed_schedule_segments = sum([convert_time_to_segment(self._START_HOUR, 
                                                                       self._END_HOUR, 
                                                                       self._TIME_UNIT, 
                                                                       fs) for fs in schedule], [])
                all_time_segments = convert_time_to_segment(self._START_HOUR, self._END_HOUR, self._TIME_UNIT)
                free_time = [s for s in range(len(all_time_segments)) if s not in fixed_schedule_segments]
                
                if len(free_time):
                    valid_time_segments = [seg for seg in group_consecutive_segments(free_time) if len(seg) >= min_time_slot_n]
                    for valid_time in valid_time_segments:
                        if (valid_time[0] < prediction_schedule_segments[0] and pred_date == date) or (len(valid_time) and compare_iso_time(pred_date, date)):
                            free_max_st, _ = convert_segment_to_time(self._START_HOUR, self._END_HOUR, self._TIME_UNIT, [valid_time[0]])
                            free_max_st_iso = get_iso_time(free_max_st, date, utc_offset=utc_offset)
                            if compare_iso_time(free_max_st_iso, current_time):
                                return False
        return True
    

    def schedule_check(self,
                      prediction: Union[str, dict], 
                      gt_patient_condition: dict,
                      doctor_information: dict,
                      environment: HospitalEnvironment) -> Tuple[bool, str]:
        """
        Validates a predicted schedule for a doctor by checking its structure, time validity, 
        duplication with existing schedules, and updates the doctor's schedule if valid.

        Args:
            prediction (Union[str, dict]): The predicted allocation result, either a string (if parsing failed)
                                           or a dictionary mapping a doctor's name to a schedule with 'start' and 'end' times.
            gt_patient_condition (dict): Ground truth patient conditions used only for sanity checks.
            doctor_information (dict): Dictionary of doctor data including their existing schedules.
                                       Each key is a doctor's name, and each value includes a 'schedule' field.
            environment (HospitalEnvironment): Hospital environment instance to manage patient schedules.

        Returns:
            Tuple[bool, str]: 
                - A boolean indicating whether the prediction passed all sanity checks.
                - A string explaining its status.
        """
        ############################ Check the prediciton format #############################
        if not isinstance(prediction, dict):
            return False, STATUS_CODES['format']    # Could not be parsed as a dictionary
        elif len(prediction['schedule']) > 1:
            return False, STATUS_CODES['conflict']['physician']    # Allocated more than one doctor; cannot determine target
    
        ################## Check the predicted schedule type and validities ##################
        try:
            pred_doctor_name = list(prediction['schedule'].keys())[0]
            start = prediction['schedule'][pred_doctor_name]['start']
            end = prediction['schedule'][pred_doctor_name]['end']
            date = prediction['schedule'][pred_doctor_name]['date']
            fixed_schedules = doctor_information[pred_doctor_name]['schedule']
            start_iso_time = get_iso_time(start, date, utc_offset=environment._utc_offset)
            assert isinstance(start, float) and isinstance(end, float) and isinstance(date, str) \
                and start < end and start >= self._START_HOUR and end <= self._END_HOUR \
                and compare_iso_time(start_iso_time, environment.current_time) and date in fixed_schedules
            assert gt_patient_condition['department'] == doctor_information[pred_doctor_name]['department']
            
            # Duration mismatched case
            if not float(Decimal(str(1)) / Decimal(str(doctor_information[pred_doctor_name]['capacity_per_hour']))) == float(Decimal(str(end)) - Decimal(str(start))):
                return False, STATUS_CODES['duration']
            
        except KeyError:
            return False, STATUS_CODES['format']    # Schedule allocation missing or doctor not found
        except AssertionError:
            return False, STATUS_CODES['schedule']    # Invalid schedule times or department

        ####################### Check the duplication of the schedules #######################
        prediction_schedule_segments = convert_time_to_segment(self._START_HOUR,
                                                               self._END_HOUR,
                                                               self._TIME_UNIT,
                                                               [start, end])
        fixed_schedule_segments = sum([convert_time_to_segment(self._START_HOUR, 
                                                               self._END_HOUR, 
                                                               self._TIME_UNIT, 
                                                               fs) for fs in fixed_schedules[date]], [])
        
        if len(set(prediction_schedule_segments) & set(fixed_schedule_segments)):
            return False, STATUS_CODES['conflict']['time']    # Overlaps with an existing schedule
        
        ####################### Check the patient's preferences  #######################
        if gt_patient_condition['preference'] == 'doctor':
            if gt_patient_condition.get('preferred_doctor') != pred_doctor_name:
                return False, STATUS_CODES['preference']['physician']
        
        if gt_patient_condition['preference'] == 'date':
            if compare_iso_time(gt_patient_condition.get('valid_from'), date):
                return False, STATUS_CODES['preference']['date']
        
        is_earliest = self.__check_is_earliest(
            prediction, 
            gt_patient_condition,
            doctor_information,
            environment,
        )

        if not is_earliest:
            return False, STATUS_CODES['preference']['asap']
        
        return True, STATUS_CODES['correct']
