MedMate-Medicine-Tracker/config_flow.py

"""Config flow for MedMate integration."""
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:
                self._medicine_data = user_input
                return await self.async_step_schedule()

        schema = vol.Schema({
            vol.Required(CONF_MEDICINE_NAME): str,
            vol.Required(CONF_ACTIVE_INGREDIENT): str,
            vol.Optional(CONF_STRENGTH, default=""): str,
            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:
            self._schedule_data = {
                CONF_DAYS: user_input.get(CONF_DAYS, []),
                CONF_TIMES: user_input.get(CONF_TIMES, [])
            }
            return await self.async_step_prescription()

        schema = vol.Schema({
            vol.Required(CONF_DAYS, default=DAYS_OF_WEEK): cv.multi_select(
                {day: day.title() for day in DAYS_OF_WEEK}
            ),
            vol.Required(CONF_TIMES, default=["morning"]): cv.multi_select(TIME_SLOTS),
        })

        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:
            self._prescription_data = user_input
            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=""): str,
            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 CONF_REPEATS_LEFT not in self._prescription_data:
            self._prescription_data[CONF_REPEATS_LEFT] = self._prescription_data[CONF_TOTAL_REPEATS]

        medicine_data = {
            **self._medicine_data,
            CONF_SCHEDULE: self._schedule_data,
            CONF_PRESCRIPTION: self._prescription_data,
            "inventory": 0,
            "last_taken": {},
        }

        # Store the 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._medicine_data[CONF_MEDICINE_NAME].lower()
        medicines[medicine_id] = medicine_data
        
        await store.async_save({"medicines": medicines})

        return self.async_create_entry(
            title=f"MedMate - {self._medicine_data[CONF_MEDICINE_NAME]}",
            data={"medicine_id": medicine_id},
        )

    @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:
            self._medicine_data = user_input
            return await self.async_step_schedule()

        schema = vol.Schema({
            vol.Required(
                CONF_MEDICINE_NAME, 
                default=current_medicine.get(CONF_MEDICINE_NAME, "")
            ): str,
            vol.Required(
                CONF_ACTIVE_INGREDIENT,
                default=current_medicine.get(CONF_ACTIVE_INGREDIENT, "")
            ): str,
            vol.Optional(
                CONF_STRENGTH, 
                default=current_medicine.get(CONF_STRENGTH, "")
            ): str,
            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, {})
        
        if user_input is not None:
            self._schedule_data = {
                CONF_DAYS: user_input.get(CONF_DAYS, []),
                CONF_TIMES: user_input.get(CONF_TIMES, [])
            }
            return await self.async_step_prescription()

        schema = vol.Schema({
            vol.Required(
                CONF_DAYS, 
                default=current_schedule.get(CONF_DAYS, DAYS_OF_WEEK)
            ): cv.multi_select({day: day.title() for day in DAYS_OF_WEEK}),
            vol.Required(
                CONF_TIMES, 
                default=current_schedule.get(CONF_TIMES, ["morning"])
            ): cv.multi_select(TIME_SLOTS),
        })

        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:
            self._prescription_data = user_input
            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, "")
            ): str,
            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 medicine data
        updated_medicine = {
            **current_medicine,
            **self._medicine_data,
            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={})