# src/tools/viz_executor.py
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime
from typing import List, Dict, Any
from domain.training.charts import ChartSpec, ChartType
from .interfaces import IVisualizationExecutor
from src.tools.runner_ai import calculate_pace_s_per_km


class TemplateVisualizationExecutor(IVisualizationExecutor):
    def render_chart(self, chart_spec: ChartSpec, features: List[Dict[str, Any]]) -> plt.Figure:
        if chart_spec.chart_type == ChartType.PACE:
            return self._render_pace_chart(features, chart_spec.title or "Pace over Time")
        elif chart_spec.chart_type == ChartType.HEART_RATE:
            return self._render_hr_chart(features, chart_spec.title or "Heart Rate over Time")
        elif chart_spec.chart_type == ChartType.VOLUME:
            return self._render_volume_chart(features, chart_spec.title or "Weekly Distance")
        elif chart_spec.chart_type == ChartType.FREQUENCY:
            return self._render_frequency_chart(features, chart_spec.title or "Run Frequency")
        elif chart_spec.chart_type == ChartType.ZONES:
            return self._render_zones_chart(features, chart_spec.title or "Effort Zones")
        elif chart_spec.chart_type == ChartType.DYNAMIC:
            return self._render_dynamic_chart(chart_spec.params, chart_spec.title or "Data Visualization")
        else:
            return self._empty_chart(f"Unsupported Chart Type: {chart_spec.chart_type}")

    def _render_pace_chart(self, features: List[Dict[str, Any]], title: str) -> plt.Figure:
        fig, ax = plt.subplots(figsize=(10, 5))
        dates = []
        paces = []

        for run in features:
            dt = run.get("start_time")
            pace_s = calculate_pace_s_per_km(
                run.get("total_distance_m"), run.get("total_duration_s")
            )
            if dt and pace_s:
                dates.append(dt)
                # Convert s/km to min/km (float)
                paces.append(pace_s / 60.0)

        if not dates:
            return self._empty_chart(title)

        # Convert strings to datetime objects if needed
        processed_dates = []
        for d in dates:
            if isinstance(d, str):
                try:
                    # Handle Z suffix for UTC
                    processed_dates.append(datetime.fromisoformat(d.replace("Z", "+00:00")))
                except ValueError:
                    processed_dates.append(d)
            else:
                processed_dates.append(d)

        ax.plot(processed_dates, paces, marker="o", linestyle="-", color="blue")
        ax.set_title(title)
        ax.set_ylabel("Pace (min/km)")
        ax.set_xlabel("Date")
        ax.grid(True)
        # Invert y-axis for pace (smaller min/km is faster)
        ax.invert_yaxis()
        
        # Format dates as DD-MMM
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%d-%b'))
        fig.autofmt_xdate()
        plt.close(fig)
        
        return fig

    def _render_hr_chart(self, features: List[Dict[str, Any]], title: str) -> plt.Figure:
        fig, ax = plt.subplots(figsize=(10, 5))
        dates = []
        hrs = []

        for run in features:
            dt = run.get("start_time")
            hr = run.get("avg_hr_bpm")
            if dt and hr:
                dates.append(dt)
                hrs.append(hr)

        if not dates:
            return self._empty_chart(title)

        # Convert strings to datetime objects if needed
        processed_dates = []
        for d in dates:
            if isinstance(d, str):
                try:
                    processed_dates.append(datetime.fromisoformat(d.replace("Z", "+00:00")))
                except ValueError:
                    processed_dates.append(d)
            else:
                processed_dates.append(d)

        ax.plot(processed_dates, hrs, marker="s", linestyle="--", color="red")
        ax.set_title(title)
        ax.set_ylabel("Heart Rate (bpm)")
        ax.set_xlabel("Date")
        ax.grid(True)
        
        # Format dates as DD-MMM
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%d-%b'))
        fig.autofmt_xdate()
        return fig

    def _render_volume_chart(self, features: List[Dict[str, Any]], title: str) -> plt.Figure:
        from collections import defaultdict
        fig, ax = plt.subplots(figsize=(10, 5))
        
        # Group by week
        weekly_km = defaultdict(float)
        for f in features:
            dt = f.get("start_time")
            if isinstance(dt, str):
                dt = datetime.fromisoformat(dt.replace("Z", "+00:00"))
            if dt:
                # ISO Week
                year, week, _ = dt.isocalendar()
                key = f"{year}-W{week:02d}"
                weekly_km[key] += f.get("total_distance_m", 0) / 1000.0
                
        if not weekly_km:
            return self._empty_chart(title)
            
        sorted_weeks = sorted(weekly_km.keys())
        km_values = [weekly_km[w] for w in sorted_weeks]
        
        ax.bar(sorted_weeks, km_values, color="teal", alpha=0.7)
        ax.set_title(title)
        ax.set_ylabel("Distance (km)")
        ax.set_xlabel("Week")
        ax.grid(axis='y', linestyle='--', alpha=0.7)
        return fig

    def _render_frequency_chart(self, features: List[Dict[str, Any]], title: str) -> plt.Figure:
        from collections import defaultdict
        fig, ax = plt.subplots(figsize=(10, 5))
        
        # Group by day of week
        dow_counts = defaultdict(int)
        days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
        
        for f in features:
            dt = f.get("start_time")
            if isinstance(dt, str):
                dt = datetime.fromisoformat(dt.replace("Z", "+00:00"))
            if dt:
                dow = dt.weekday()
                dow_counts[dow] += 1
                
        if not dow_counts:
            return self._empty_chart(title)
            
        counts = [dow_counts[i] for i in range(7)]
        
        ax.bar(days, counts, color="orange", alpha=0.7)
        ax.set_title(title)
        ax.set_ylabel("Number of Runs")
        ax.grid(axis='y', linestyle='--', alpha=0.7)
        return fig

    def _render_zones_chart(self, features: List[Dict[str, Any]], title: str) -> plt.Figure:
        # Simple placeholder for HR zones distribution if data available
        # In this MVP, we might not have zones calculated per run, but we can do a mock distribution
        fig, ax = plt.subplots(figsize=(8, 8))
        labels = ['Zone 1', 'Zone 2', 'Zone 3', 'Zone 4', 'Zone 5']
        
        # Mock logic based on average HR
        avg_hrs = [f.get("avg_hr_bpm", 0) for f in features if f.get("avg_hr_bpm")]
        if not avg_hrs:
            return self._empty_chart(title)
            
        # Simplified: just count runs whose avg HR falls into zones (assuming max HR 190)
        distribution = [0] * 5
        for hr in avg_hrs:
            if hr < 130: distribution[0] += 1
            elif hr < 150: distribution[1] += 1
            elif hr < 165: distribution[2] += 1
            elif hr < 180: distribution[3] += 1
            else: distribution[4] += 1
            
        ax.pie(distribution, labels=labels, autopct='%1.1f%%', colors=plt.cm.RdYlGn_r([0.2, 0.4, 0.6, 0.8, 1.0]))
        ax.set_title(title)
        return fig

    def _render_dynamic_chart(self, params: Dict[str, Any], title: str) -> plt.Figure:
        labels = params.get("labels", [])
        values = params.get("values", [])
        chart_type = params.get("type", "bar")
        
        if not labels or not values:
            return self._empty_chart(title)
            
        fig, ax = plt.subplots(figsize=(10, 5))
        if chart_type == "bar":
            ax.bar(labels, values, color="skyblue")
        elif chart_type == "pie":
            ax.pie(values, labels=labels, autopct='%1.1f%%')
        else:
            ax.plot(labels, values, marker='o')
            
        ax.set_title(title)
        ax.grid(True, axis='y', alpha=0.3)
        return fig

    def _empty_chart(self, title: str) -> plt.Figure:
        fig, ax = plt.subplots(figsize=(6, 3))
        ax.text(0.5, 0.5, "No Data Available", ha="center", va="center")
        ax.set_title(title)
        ax.axis("off")
        return fig