import json
import panel
from datetime import datetime, timedelta

from pydantic import BaseModel

from watt42_viewlib import attach_w42_state

w42_state = panel.rx(None)
attach_w42_state(rx_var=w42_state, system_id="79476e53-dea6-44fa-976c-eff6260baeb6")

state_as_text = panel.bind(lambda s: f"W42 State:\n\n```\n{json.dumps(s, indent=2)}\n```\n\nReplace this with some awesome visuals", w42_state)

class Output(BaseModel):
    target_storage: float
    target_geyser_temperature: float
    storage: list[float]
    grid_sell: list[float]
    grid_buy: list[float]


class PvForecast(BaseModel):
    at: datetime = datetime.now()
    slots: list[float] = []

class LoadForecast(BaseModel):
    at: datetime = datetime.now()
    slots: list[float] = []

class SystemState(BaseModel):
    load_forecast: LoadForecast = LoadForecast()
    pv_forecast: PvForecast = PvForecast()
    now: datetime = datetime.now()
    sell_prices: list[float] = []
    buy_prices: list[float] = []
    output: Output | None = None

state: SystemState = panel.rx(lambda s: SystemState.model_validate(s) if s else SystemState())(w42_state)

state_pane = panel.pane.Markdown(state_as_text, sizing_mode='stretch_width')

def now_from_state(state: SystemState) -> datetime:
    return state.load_forecast.at

# now_from_state_rx = panel.rx(now_from_state)(state)

def get_label(at: datetime, slot_index: int) -> str:
    slot_time = at + timedelta(minutes=15 * slot_index)
    return slot_time.strftime('%H:%M')

def load_fc_chart(state: SystemState) -> dict:
    slots = state.load_forecast.slots
    at = state.load_forecast.at
    storage = state.output.storage if state.output else [0] * len(slots)
    battery = [storage[ix] - storage[ix-1] if ix > 0 else 0 for ix in range(len(slots))] if state.output else [0] * len(slots)
    grid = [state.output.grid_buy[ix] - state.output.grid_sell[ix] if state.output else 0 for ix in range(len(slots))]
    return {
        'title': {
            'text': '24hr Forecast'
        },
        'legend': {
            'data': ['Load Forecast', 'Battery Storage', 'Battery in/out', 'PV Forecast', 'Grid in/out']
        },
        'tooltip': {
            'trigger': 'axis'
        },
        'xAxis': {
            'type': 'category',
            'data': [get_label(at, i) for i in range(len(slots))]
        },
        'yAxis': {
            'type': 'value'
        },
        'series': [
            {
                'name': 'Load Forecast',
                'data': slots,
                'type': 'line'
            }, {
                'name': 'Battery Storage',
                'data': storage,
                'type': 'line',
                'color': 'orange'
            }, {
                'name': 'Battery in/out',
                'data': battery,
                'type': 'line',
                'color': 'lightblue'
            }, {
				'name': 'PV Forecast',
				'data': state.pv_forecast.slots,
                'type': 'line',
                'color': 'green'
			}, {
			    'name': 'Grid in/out',
                'data': grid,
                'type': 'line',
                'color': 'red'
			}
        ]
    }

load_fc_chart_rx = panel.rx(load_fc_chart)(state)

def prices_chart(state: SystemState) -> dict:
    return {
        'title': {
            'text': 'Electricity Prices'
        },
        'legend': {
            'data': ['Sell Prices', 'Buy Prices']
        },
        'tooltip': {
            'trigger': 'axis'
        },
        'xAxis': {
            'type': 'category',
            'data': [get_label(state.now, i) for i in range(len(state.sell_prices))]
        },
        'yAxis': {
            'type': 'value'
        },
        'series': [
            {
                'name': 'Sell Prices',
                'data': state.sell_prices,
                'type': 'line',
                'color': 'green'
            }, {
                'name': 'Buy Prices',
                'data': state.buy_prices,
                'type': 'line',
                'color': 'red'
            }
        ]
    }

prices_chart_rx = panel.rx(prices_chart)(state)

datetime_fmt = "%Y-%m-%d %H:%M:%S"

value = w42_state.rx.value

_ = panel.template.FastListTemplate(
    title="Sample W42 View App",
    sidebar=[panel.pane.Markdown("This is a sample sidebar.")],
    main=[
        # panel.pane.Markdown(state_as_text, sizing_mode='stretch_width'),
        panel.pane.ECharts(
            load_fc_chart_rx,
            sizing_mode='stretch_width',
            height=400
        ),
        panel.pane.ECharts(
            prices_chart_rx,
            sizing_mode='stretch_width',
            height=300
        ),
        w42_state,
        panel.pane.Markdown(panel.bind(lambda s: f"State at {s.now.strftime(datetime_fmt)}, Load Forecast Time: {s.load_forecast.at}", state), sizing_mode='stretch_width'),
    ],
).servable()