MedMate-Medicine-Tracker/config_flow.py

"""Config flow for MedMate integration - Bulletproof version."""
from __future__ import annotations

import logging
from typing import Any
import voluptuous as vol
from datetime import datetime, date

from homeassistant import config_entries
from homeassistant.core import HomeAssistant, callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers import config_validation as cv
from homeassistant.helpers.storage import Store

from .const import (
    DOMAIN,
    CONF_MEDICINE_NAME,
    CONF_ACTIVE_INGREDIENT,
    CONF_STRENGTH,
    CONF_PACK_SIZE,
    CONF_SCHEDULE,
    CONF_PRESCRIPTION,
    CONF_DAYS,
    CONF_TIMES,
    CONF_ISSUE_DATE,
    CONF_EXPIRY_DATE,
    CONF_DOCTOR,
    CONF_TOTAL_REPEATS,
    CONF_REPEATS_LEFT,
    TIME_SLOTS,
    DAYS_OF_WEEK,
    DEFAULT_PACK_SIZE,
    DEFAULT_REPEATS,
    STORAGE_KEY,
    STORAGE_VERSION,
)

_LOGGER = logging.getLogger(__name__)


class MedMateConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
    """Handle a config flow for MedMate."""

    VERSION = 1

    def __init__(self) -> None:
        """Initialize the config flow."""
        self._medicine_data = {}
        self._schedule_data = {}
        self._prescription_data = {}
        self._editing_medicine_id = None
        
    async def async_step_user(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle the initial step."""
        if user_input is not None:
            return await self.async_step_medicine_basic()

        return self.async_show_form(
            step_id="user",
            data_schema=vol.Schema({}),
            description_placeholders={
                "name": "MedMate - Medicine Tracker"
            }
        )

    async def async_step_medicine_basic(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle basic medicine information."""
        errors = {}

        if user_input is not None:
            # Validate medicine name is unique
            store = Store(self.hass, STORAGE_VERSION, STORAGE_KEY)
            existing_data = await store.async_load() or {}
            medicines = existing_data.get("medicines", {})
            
            medicine_name = user_input[CONF_MEDICINE_NAME].lower()
            if medicine_name in medicines and self._editing_medicine_id != medicine_name:
                errors[CONF_MEDICINE_NAME] = "name_exists"
            else:
                # Store ONLY the expected keys - CRITICAL FILTERING
                self._medicine_data = {
                    CONF_MEDICINE_NAME: user_input[CONF_MEDICINE_NAME],
                    CONF_ACTIVE_INGREDIENT: user_input[CONF_ACTIVE_INGREDIENT],
                    CONF_STRENGTH: user_input.get(CONF_STRENGTH, ""),
                    CONF_PACK_SIZE: user_input.get(CONF_PACK_SIZE, DEFAULT_PACK_SIZE),
                }
                return await self.async_step_schedule()

        schema = vol.Schema({
            vol.Required(CONF_MEDICINE_NAME): cv.string,
            vol.Required(CONF_ACTIVE_INGREDIENT): cv.string,
            vol.Optional(CONF_STRENGTH, default=""): cv.string,
            vol.Optional(CONF_PACK_SIZE, default=DEFAULT_PACK_SIZE): cv.positive_int,
        })

        return self.async_show_form(
            step_id="medicine_basic",
            data_schema=schema,
            errors=errors,
        )

    async def async_step_schedule(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle schedule configuration."""
        if user_input is not None:
            # CRITICAL: Extract only the actual values, ignore ALL form keys
            selected_days = []
            selected_times = []
            
            # Process days - extract values from checkbox form keys
            for day in DAYS_OF_WEEK:
                form_key = f"day_{day}"
                if user_input.get(form_key, False):
                    selected_days.append(day)
                    
            # Process times - extract values from checkbox form keys
            for time_key in TIME_SLOTS.keys():
                form_key = f"time_{time_key}"
                if user_input.get(form_key, False):
                    selected_times.append(time_key)
            
            # Store ONLY the clean, expected data structure
            self._schedule_data = {
                CONF_DAYS: selected_days,
                CONF_TIMES: selected_times
            }
            
            _LOGGER.debug("Processed schedule data: %s", self._schedule_data)
            return await self.async_step_prescription()

        # Create form schema with temporary checkbox keys
        # These keys are ONLY for the form, never stored permanently
        schema_fields = {}
        
        # Add day checkboxes (all selected by default)
        for day in DAYS_OF_WEEK:
            schema_fields[vol.Optional(f"day_{day}", default=True)] = cv.boolean
            
        # Add time checkboxes (only morning selected by default)
        for time_key in TIME_SLOTS.keys():
            default_val = (time_key == "morning")
            schema_fields[vol.Optional(f"time_{time_key}", default=default_val)] = cv.boolean

        schema = vol.Schema(schema_fields)

        return self.async_show_form(
            step_id="schedule",
            data_schema=schema,
        )

    async def async_step_prescription(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle prescription information."""
        if user_input is not None:
            # Store ONLY expected prescription keys
            self._prescription_data = {
                CONF_ISSUE_DATE: user_input.get(CONF_ISSUE_DATE),
                CONF_EXPIRY_DATE: user_input.get(CONF_EXPIRY_DATE),
                CONF_DOCTOR: user_input.get(CONF_DOCTOR, ""),
                CONF_TOTAL_REPEATS: user_input.get(CONF_TOTAL_REPEATS, DEFAULT_REPEATS),
                CONF_REPEATS_LEFT: user_input.get(CONF_REPEATS_LEFT),
            }
            return await self.async_create_entry_data()

        today = date.today()
        expiry_date = date(today.year + 1, today.month, today.day)

        schema = vol.Schema({
            vol.Optional(CONF_ISSUE_DATE, default=today): cv.date,
            vol.Optional(CONF_EXPIRY_DATE, default=expiry_date): cv.date,
            vol.Optional(CONF_DOCTOR, default=""): cv.string,
            vol.Optional(CONF_TOTAL_REPEATS, default=DEFAULT_REPEATS): cv.positive_int,
            vol.Optional(CONF_REPEATS_LEFT): cv.positive_int,
        })

        return self.async_show_form(
            step_id="prescription",
            data_schema=schema,
        )

    async def async_create_entry_data(self) -> FlowResult:
        """Create the config entry."""
        # Set repeats_left to total_repeats if not specified
        if not self._prescription_data.get(CONF_REPEATS_LEFT):
            self._prescription_data[CONF_REPEATS_LEFT] = self._prescription_data[CONF_TOTAL_REPEATS]

        # Build the complete medicine data with ONLY expected keys
        medicine_data = {
            CONF_MEDICINE_NAME: self._medicine_data[CONF_MEDICINE_NAME],
            CONF_ACTIVE_INGREDIENT: self._medicine_data[CONF_ACTIVE_INGREDIENT],
            CONF_STRENGTH: self._medicine_data[CONF_STRENGTH],
            CONF_PACK_SIZE: self._medicine_data[CONF_PACK_SIZE],
            CONF_SCHEDULE: self._schedule_data,  # Contains only CONF_DAYS and CONF_TIMES
            CONF_PRESCRIPTION: self._prescription_data,
            "inventory": 0,
            "last_taken": {},
        }

        # Store the medicine data in the persistent storage
        store = Store(self.hass, STORAGE_VERSION, STORAGE_KEY)
        existing_data = await store.async_load() or {}
        medicines = existing_data.get("medicines", {})
        
        medicine_id = self._medicine_data[CONF_MEDICINE_NAME].lower()
        medicines[medicine_id] = medicine_data
        
        await store.async_save({"medicines": medicines})

        # Create config entry with MINIMAL data - only what Home Assistant needs
        # CRITICAL: Only store the medicine_id, nothing else that could cause validation issues
        config_entry_data = {
            "medicine_id": medicine_id
        }
        
        _LOGGER.debug("Creating config entry with data: %s", config_entry_data)

        return self.async_create_entry(
            title=f"MedMate - {self._medicine_data[CONF_MEDICINE_NAME]}",
            data=config_entry_data,  # Only the minimal required data
        )

    @staticmethod
    @callback
    def async_get_options_flow(
        config_entry: config_entries.ConfigEntry,
    ) -> MedMateOptionsFlow:
        """Get the options flow for this handler."""
        return MedMateOptionsFlow(config_entry)


class MedMateOptionsFlow(config_entries.OptionsFlow):
    """Handle options flow for MedMate."""

    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
        """Initialize options flow."""
        self.config_entry = config_entry
        self._medicine_data = {}
        self._schedule_data = {}
        self._prescription_data = {}

    async def async_step_init(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Manage the options."""
        return await self.async_step_medicine_basic()

    async def async_step_medicine_basic(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle basic medicine information for options."""
        errors = {}
        
        # Load current medicine data
        store = Store(self.hass, STORAGE_VERSION, STORAGE_KEY)
        existing_data = await store.async_load() or {}
        medicines = existing_data.get("medicines", {})
        medicine_id = self.config_entry.data.get("medicine_id")
        current_medicine = medicines.get(medicine_id, {})

        if user_input is not None:
            # Filter to only expected keys
            self._medicine_data = {
                CONF_MEDICINE_NAME: user_input[CONF_MEDICINE_NAME],
                CONF_ACTIVE_INGREDIENT: user_input[CONF_ACTIVE_INGREDIENT],
                CONF_STRENGTH: user_input.get(CONF_STRENGTH, ""),
                CONF_PACK_SIZE: user_input.get(CONF_PACK_SIZE, DEFAULT_PACK_SIZE),
            }
            return await self.async_step_schedule()

        schema = vol.Schema({
            vol.Required(
                CONF_MEDICINE_NAME, 
                default=current_medicine.get(CONF_MEDICINE_NAME, "")
            ): cv.string,
            vol.Required(
                CONF_ACTIVE_INGREDIENT,
                default=current_medicine.get(CONF_ACTIVE_INGREDIENT, "")
            ): cv.string,
            vol.Optional(
                CONF_STRENGTH, 
                default=current_medicine.get(CONF_STRENGTH, "")
            ): cv.string,
            vol.Optional(
                CONF_PACK_SIZE, 
                default=current_medicine.get(CONF_PACK_SIZE, DEFAULT_PACK_SIZE)
            ): cv.positive_int,
        })

        return self.async_show_form(
            step_id="medicine_basic",
            data_schema=schema,
            errors=errors,
        )

    async def async_step_schedule(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle schedule configuration for options."""
        # Load current medicine data
        store = Store(self.hass, STORAGE_VERSION, STORAGE_KEY)
        existing_data = await store.async_load() or {}
        medicines = existing_data.get("medicines", {})
        medicine_id = self.config_entry.data.get("medicine_id")
        current_medicine = medicines.get(medicine_id, {})
        current_schedule = current_medicine.get(CONF_SCHEDULE, {})
        current_days = current_schedule.get(CONF_DAYS, DAYS_OF_WEEK)
        current_times = current_schedule.get(CONF_TIMES, ["morning"])
        
        if user_input is not None:
            # Extract clean values from form checkboxes
            selected_days = []
            selected_times = []
            
            # Process days
            for day in DAYS_OF_WEEK:
                if user_input.get(f"day_{day}", False):
                    selected_days.append(day)
                    
            # Process times  
            for time_key in TIME_SLOTS.keys():
                if user_input.get(f"time_{time_key}", False):
                    selected_times.append(time_key)
                    
            # Store only clean data
            self._schedule_data = {
                CONF_DAYS: selected_days,
                CONF_TIMES: selected_times
            }
            return await self.async_step_prescription()

        # Create form schema with checkboxes
        schema_fields = {}
        
        # Add day checkboxes with current values as defaults
        for day in DAYS_OF_WEEK:
            default_val = day in current_days
            schema_fields[vol.Optional(f"day_{day}", default=default_val)] = cv.boolean
            
        # Add time checkboxes with current values as defaults
        for time_key in TIME_SLOTS.keys():
            default_val = time_key in current_times
            schema_fields[vol.Optional(f"time_{time_key}", default=default_val)] = cv.boolean

        schema = vol.Schema(schema_fields)

        return self.async_show_form(
            step_id="schedule",
            data_schema=schema,
        )

    async def async_step_prescription(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle prescription information for options."""
        # Load current medicine data
        store = Store(self.hass, STORAGE_VERSION, STORAGE_KEY)
        existing_data = await store.async_load() or {}
        medicines = existing_data.get("medicines", {})
        medicine_id = self.config_entry.data.get("medicine_id")
        current_medicine = medicines.get(medicine_id, {})
        current_prescription = current_medicine.get(CONF_PRESCRIPTION, {})
        
        if user_input is not None:
            # Filter to expected prescription keys only
            self._prescription_data = {
                CONF_ISSUE_DATE: user_input.get(CONF_ISSUE_DATE),
                CONF_EXPIRY_DATE: user_input.get(CONF_EXPIRY_DATE),
                CONF_DOCTOR: user_input.get(CONF_DOCTOR, ""),
                CONF_TOTAL_REPEATS: user_input.get(CONF_TOTAL_REPEATS, DEFAULT_REPEATS),
                CONF_REPEATS_LEFT: user_input.get(CONF_REPEATS_LEFT, DEFAULT_REPEATS),
            }
            return await self.async_update_entry_data()

        today = date.today()
        
        schema = vol.Schema({
            vol.Optional(
                CONF_ISSUE_DATE,
                default=current_prescription.get(CONF_ISSUE_DATE, today)
            ): cv.date,
            vol.Optional(
                CONF_EXPIRY_DATE,
                default=current_prescription.get(CONF_EXPIRY_DATE, date(today.year + 1, today.month, today.day))
            ): cv.date,
            vol.Optional(
                CONF_DOCTOR,
                default=current_prescription.get(CONF_DOCTOR, "")
            ): cv.string,
            vol.Optional(
                CONF_TOTAL_REPEATS,
                default=current_prescription.get(CONF_TOTAL_REPEATS, DEFAULT_REPEATS)
            ): cv.positive_int,
            vol.Optional(
                CONF_REPEATS_LEFT,
                default=current_prescription.get(CONF_REPEATS_LEFT, DEFAULT_REPEATS)
            ): cv.positive_int,
        })

        return self.async_show_form(
            step_id="prescription",
            data_schema=schema,
        )

    async def async_update_entry_data(self) -> FlowResult:
        """Update the config entry."""
        # Load and update medicine data
        store = Store(self.hass, STORAGE_VERSION, STORAGE_KEY)
        existing_data = await store.async_load() or {}
        medicines = existing_data.get("medicines", {})
        
        medicine_id = self.config_entry.data.get("medicine_id")
        current_medicine = medicines.get(medicine_id, {})
        
        # Update with only expected keys
        updated_medicine = {
            **current_medicine,
            CONF_MEDICINE_NAME: self._medicine_data[CONF_MEDICINE_NAME],
            CONF_ACTIVE_INGREDIENT: self._medicine_data[CONF_ACTIVE_INGREDIENT],
            CONF_STRENGTH: self._medicine_data[CONF_STRENGTH],
            CONF_PACK_SIZE: self._medicine_data[CONF_PACK_SIZE],
            CONF_SCHEDULE: self._schedule_data,
            CONF_PRESCRIPTION: self._prescription_data,
        }
        
        medicines[medicine_id] = updated_medicine
        await store.async_save({"medicines": medicines})

        return self.async_create_entry(title="", data={})