from core.pipeline.step import PipelineStep
from domain.training.charts import ChartSpec, ChartType
from tools.helpers import serialize_charts

DEFAULT_CHART_SPECS = [
    ChartSpec(chart_type=ChartType.PACE, title="Pace"),
    ChartSpec(chart_type=ChartType.HEART_RATE, title="Heart Rate"),
]

class VisualizationStep(PipelineStep):
    name = "visualization_agent"

    def __init__(self, agent, viz_executor):
        self.agent = agent
        self.viz_executor = viz_executor

    async def run(self, context):
        import logging
        logger = logging.getLogger(__name__)
        
        # Explicit scope selection
        if context.visualization_scope == "full":
            runs = context.runs
        else:
            runs = context.week_runs
            
        logger.info(f"[VISUALISATION DEBUG] scope: {context.visualization_scope}, runs used: {len(runs)}")
        detailed_features = [
            f.model_dump(mode="json") if hasattr(f, "model_dump") else f for f in runs
        ]
        
        charts = {}
        logger.info(f"[VISUALISATION DEBUG] features count: {len(detailed_features)}")
        if detailed_features:
            # Deterministic rendering via viz_executor for dashboard
            for spec in DEFAULT_CHART_SPECS:
                fig = self.viz_executor.render_chart(spec, detailed_features)
                logger.info(f"[VISUALISATION DEBUG] chart generated: {fig is not None}")
                if fig:
                    if spec.chart_type == ChartType.PACE:
                        charts["pace_chart"] = fig
                    elif spec.chart_type == ChartType.HEART_RATE:
                        charts["hr_chart"] = fig
                    else:
                        key = spec.title or f"chart_{spec.chart_type}"
                        charts[key] = fig
        context.charts = charts