Proper Game State and UI Management

I am having the hardest time wrapping my head around a proper game state and UI management system.

So lets say I have a game that has a main menu state, play state, and a game over state. Each state also has a respective UI to go along with it, as well as additional UI (menu, paused, etc..) that each can make use of if needed. Each state also has its own behavior. The main menu with be listening for inputs, maybe playing some animations, and other things. The play state will be controlling everything that happens during game play (most likely separated into many different classes being "managed" by this state). This includes score, money, player interactions, etc... The game over state is responsible for game over stuff... I think you get the point.

Previously I would manage this the following way: I would have a parent game object called StateController that had child objects for MainMenuState, PlayState, and GameOverState. Each of the children would have their respective state scripts AND UI elements/scripts. The StateController would be responsible for calling SetActive(true) on whatever state the game was currently in, and SetActive(false) for whatever the previous state was, usually controlled by my GameManager singleton.

So this worked, but now I am attempting a project of a much larger scale and I really have the feeling I should be separating these into two separate entities. My current plan is to have a UIManager game object/script that has each respective UI underneath it; MainMenuUI, PlayUI, GameOverUI. Similarly, I would have a StateManager that acts almost the same but rather than UI elements as children I would have the scripts needed for each state: MainMenuState, PlayState, and GameOverState. These would work together by having the StateManager keep a reference to UIManager and turn on/off the UI as needed.

So this was all going according to plan UNTIL I realized how closely these systems need to work together. The states constantly need to be updating and telling the UI what to display and when to display it. Trying to be as clean as possible, I want to avoid having direct references to any specific UI scripts within my state scripts (I always avoid assigning anything in the editor). My idea was to do something such as have the PlayState get the PlayUI from the UIManager and then update things from there. Or inversely have the PlayUI get the PlayState script and pull the data that it needs. But I am having the hardest time deciding on a solution and keeping it.

My question is are either of these solutions viable/reasonable? Is it smarter to go with the second option? Or are both of these solutions terrible and I am missing something huge?

My apologies for the long post, I appreciate anyone who took the time to read it.