import src.prompts.prompt_templates as prompt_templates
from src import config
from src.prompts.fieldset import FieldSet
import src.prompts.prompt_templates as prompt_templates
from constants.constants import Constants
import logging

## Carveout questions
MULTIPLE_PROCEDURE_REDUCTION_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="multiple_procedure_reductions"
).print_prompt_dict()
NEVER_EVENT_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="never_event"
).print_prompt_dict()
EXPERIMENTAL_INVESTIGATIONAL_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="experimental_investigational"
).print_prompt_dict()
# no questions for not_medically_necessary, global_services, bundled_codes, midlevels,
# technical_component, professional_component, diagnostic_procedures, preadmission,
# post_discharge, or readmissions

## Special case questions
STOP_LOSS_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="stop_loss"
).print_prompt_dict()
OUTLIER_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="outlier"
).print_prompt_dict()
SEQUESTRATION_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="sequestration"
).print_prompt_dict()
DISCOUNT_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="discount"
).print_prompt_dict()
PREMIUM_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="premium"
).print_prompt_dict()
ADDITION_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="addition"
).print_prompt_dict()
TRIGGER_CAP_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="trigger_cap"
).print_prompt_dict()
RATE_ESCALATOR_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="rate_escalator"
).print_prompt_dict()
GROUPER_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="grouper"
).print_prompt_dict()


CARVEOUT_CONFIGS = {
    # Carveouts
    "BASE_COVERED_SERVICES": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": None,
    },
    "PAYMENT_CARVEOUT": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "UNLISTED_CODE": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": True,
        "term_key": None,
    },
    "NEVER_EVENT": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "EXPERIMENTAL_INVESTIGATIONAL": {
        "questions": EXPERIMENTAL_INVESTIGATIONAL_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "NOT_MEDICALLY_NECESSARY": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "GLOBAL_SERVICES": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "BUNDLED_CODES": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "MULTIPLE_PROCEDURE_REDUCTIONS": {
        "questions": MULTIPLE_PROCEDURE_REDUCTION_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "MIDLEVELS": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "TECHNICAL_COMPONENT": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "PROFESSIONAL_COMPONENT": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "DIAGNOSTIC_PROCEDURES": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "PREADMISSION": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "POST_DISCHARGE": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "READMISSIONS": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": True,
        "default_ind": False,
        "term_key": None,
    },
    "DEFAULT_TERM": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": True,
        "term_key": None,
    },
}

SPECIAL_CASE_CONFIGS = {
    # Special Cases
    "STOP_LOSS": {
        "questions": STOP_LOSS_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "service and methodology",
        "use_two_stage": False,  # Legacy processing - this case does NOT gets its own breakout separate from the main methodology breakout
    },
    "OUTLIER": {
        "questions": OUTLIER_QUESTIONS,
        "methodology_prompt_template": prompt_templates.OUTLIER_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "service and methodology",
        "use_two_stage": True,  # Indicates special two-stage processing in which this case gets its own breakout separate from the main methodology breakout
    },
    "SEQUESTRATION": {
        "questions": SEQUESTRATION_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",
        "use_two_stage": False,  # Legacy processing - this case does NOT gets its own breakout separate from the main methodology breakout
    },
    "DISCOUNT": {
        "questions": DISCOUNT_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",
        "use_two_stage": False,  # Legacy processing - this case does NOT gets its own breakout separate from the main methodology breakout
    },
    "PREMIUM": {
        "questions": PREMIUM_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",
        "use_two_stage": False,  # Legacy processing - this case does NOT gets its own breakout separate from the main methodology breakout
    },
    "ADDITION": {
        "questions": ADDITION_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",
        "use_two_stage": False,  # Legacy processing - this case does NOT gets its own breakout separate from the main methodology breakout
    },
    "TRIGGER_CAP": {
        "questions": TRIGGER_CAP_QUESTIONS,  # Note that with two-stage processing, this case does not use the main methodology breakout questions
        "methodology_prompt_template": prompt_templates.TRIGGER_CAP_BREAKOUT,  # Use custom prompt
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",
        "use_two_stage": False,
    },
    "RATE_ESCALATOR": {
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",
        "use_two_stage": False,  # Legacy processing - this case does NOT gets its own breakout separate from the main methodology breakout
    },
    "GROUPER": {
        "questions": GROUPER_QUESTIONS,
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "service and methodology",
        "use_two_stage": True,
    },
}


def merge_special_case_configs(special_cases: list[str], constants: Constants) -> dict:
    # Start with defaults
    merged_config = {
        "questions": FieldSet(
            file_path=config.FIELD_JSON_PATH, field_type="methodology_breakout"
        ).print_prompt_dict(constants),
        "fs_questions": FieldSet(
            file_path=config.FIELD_JSON_PATH, field_type="fee_schedule_breakout"
        ).print_prompt_dict(constants),
        "methodology_prompt_template": prompt_templates.REIMB_TERM_BREAKOUT,
        "fee_schedule_prompt_template": prompt_templates.FEE_SCHEDULE_BREAKOUT,
        "carveout_ind": False,
        "default_ind": False,
        "term_key": "methodology",  # default term key
    }

    # Track which cases need "service and methodology"
    needs_service_and_methodology = []
    for case in special_cases:
        case = case.upper()
        if case in SPECIAL_CASE_CONFIGS:
            case_config = SPECIAL_CASE_CONFIGS[case]

            # Add case-specific questions (avoiding duplicates)
            case_questions = case_config.get("questions")
            if case_questions:
                if case_questions not in merged_config["questions"]:
                    merged_config["questions"] += case_questions

            # Handle custom prompt templates (last one wins, or can implement priority down the line)
            if (
                case_config.get("methodology_prompt_template")
                != prompt_templates.REIMB_TERM_BREAKOUT
            ):
                # If we get here, the template exists and is different
                merged_config["methodology_prompt_template"] = case_config[
                    "methodology_prompt_template"
                ]
                logging.info(f"Using custom methodology prompt template for {case}")
            else:
                logging.debug(
                    f"SKIPPING template assignment for {case} (using default)"
                )

            # Track term key requirements
            if case_config.get("term_key") == "service and methodology":
                needs_service_and_methodology.append(case)

        # Set term_key based on requirements
    if needs_service_and_methodology:
        merged_config["term_key"] = "service and methodology"
        logging.info(
            f"Setting term_key to 'service and methodology' due to cases: {needs_service_and_methodology}"
        )
    return merged_config