from typing import Dict, Any, List, Optional
import logging
import time
from observability import logger as obs_logger
from observability import components as obs_components
from .base import BaseAgent
from domain.training.agent_models import RiskAssessment, WeeklySummary
from domain.runner.profile import RunnerProfile

logger = logging.getLogger(__name__)


class InjuryFatigueGuardrailAgent(BaseAgent):
    """
    Heuristic-based agent that assesses the risk of injury or fatigue.
    """

    def assess_risk(
        self,
        features: WeeklySummary,
        pain_reported: bool = False,
        profile: Optional[RunnerProfile] = None,
    ) -> RiskAssessment:
        """
        Calculates risk level and reasons based on running features.
        """
        risk_level = "LOW"
        reasons = []
        adjustments = []
        metrics = {}

        # 1. Mileage Spike Heuristic
        # summary often contains 'weekly_km'
        weekly_km = features.weekly_km
        if weekly_km:
            weeks = sorted(weekly_km.items(), reverse=True)
            if len(weeks) >= 2:
                current_week_val = weeks[0][1]
                prev_week_val = weeks[1][1]
                metrics["current_week_km"] = current_week_val
                metrics["prev_week_km"] = prev_week_val

                if prev_week_val > 0:
                    increase_pct = (current_week_val - prev_week_val) / prev_week_val
                    metrics["wow_increase_pct"] = increase_pct

                    if increase_pct > 0.5:
                        risk_level = "HIGH"
                        reasons.append(f"Major mileage spike: {increase_pct:.1%} increase Wow.")
                        adjustments.append("Reduce volume significantly; include extra rest day.")
                    elif increase_pct > 0.3:
                        if risk_level != "HIGH":
                            risk_level = "MEDIUM"
                        reasons.append(
                            f"Significant mileage increase: {increase_pct:.1%} increase WoW."
                        )
                        adjustments.append("Cap intensity; avoid two hard days in a row.")

        # 1.5. Baseline Mileage Spike Heuristic
        if (
            profile
            and isinstance(profile.baseline_weekly_km, (int, float))
            and profile.baseline_weekly_km > 0
        ):
            current_week_val = 0.0
            if weekly_km:
                weeks = sorted(weekly_km.items(), reverse=True)
                if weeks:
                    current_week_val = weeks[0][1]

            if current_week_val > 0:
                baseline_increase_pct = (
                    current_week_val - profile.baseline_weekly_km
                ) / profile.baseline_weekly_km
                metrics["baseline_increase_pct"] = baseline_increase_pct

                if baseline_increase_pct > 0.6:
                    if risk_level != "HIGH":
                        risk_level = "HIGH"
                    reasons.append(
                        f"Major spike vs your baseline: {baseline_increase_pct:.1%} above baseline of {profile.baseline_weekly_km}km."
                    )
                    if "Reduce volume" not in " ".join(adjustments):
                        adjustments.append("Reduce volume to closer to your baseline.")
                elif baseline_increase_pct > 0.4:
                    if risk_level == "LOW":
                        risk_level = "MEDIUM"
                    reasons.append(
                        f"Significant increase vs your baseline: {baseline_increase_pct:.1%} above baseline."
                    )
                    if "Cap intensity" not in " ".join(adjustments):
                        adjustments.append(
                            "Be cautious with this increase; ensure adequate recovery."
                        )

        # 2. Consistency / Monotony Heuristic
        # Feature Engineering agent provides consistency_score (0-100)
        consistency = features.consistency_score
        metrics["consistency_score"] = consistency
        if consistency < 30:
            if risk_level == "LOW":
                risk_level = "MEDIUM"
            reasons.append(
                f"Low consistency score ({consistency}/100) increases injury risk during ramp-up."
            )
            adjustments.append("Build base gradually before adding intensity.")

        # 3. Pain Reported Flag (from User or scenario)
        metrics["pain_reported"] = pain_reported
        if pain_reported:
            risk_level = "HIGH"
            reasons.append("User reported pain or discomfort.")
            adjustments.append("Switch all remaining sessions this week to rest or cross-training.")

        # 4. Fatigue markers (placeholder - heart rate drift trend if provided)
        # In this PoC, we might have pace_trend from trends but it's simple.

        return RiskAssessment(
            risk_level=risk_level,
            reasons=reasons,
            recommended_adjustments=adjustments,
            metrics_used=metrics,
        )

    def run(
        self,
        features: WeeklySummary,
        pain_reported: bool = False,
        profile: Optional[RunnerProfile] = None,
    ) -> RiskAssessment:
        with obs_logger.start_span("guardrail_agent.run", obs_components.AGENT):
            return self.assess_risk(features, pain_reported, profile=profile)