# tests/simulate_lifecycle.py
import json
import datetime
import random
from typing import Dict
from unittest.mock import MagicMock

# Import the logic
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.services.rl_engine import (
    get_behavior_signal,
    update_reward,
    _load_bandit_params,
    _calculate_weekly_strain
)
from app.models.progress import Progress

def simulate_archetype(name: str, days: int, behavior_fn):
    print(f"\n{'='*40}")
    print(f" Simulating Archetype: {name} ({days} Days)")
    print(f"{'='*40}")

    user = MagicMock()
    user.user_id = f"sim_user_{name}"
    user.competence = 0.7
    user.fatigue = 0.0
    user.mood = 0.5
    user.sleep_quality = 0.8
    user.stress = 0.2
    user.current_streak = 0
    user.rl_bandit_params = None
    user.rl_history_count = 0
    user.days_since_last_login = 0
    
    # We mock the DB and _compute_historical_features to simulate a rolling 7-day progress.
    progress_history = []

    def mock_db_query(*args, **kwargs):
        mock_q = MagicMock()
        def filter_fn(*f_args, **f_kwargs):
            return mock_q
        mock_q.filter = filter_fn
        mock_q.order_by = lambda x: mock_q
        
        # We need to return exactly the mock progress history
        def all_fn():
            return progress_history
        mock_q.all = all_fn
        return mock_q

    db = MagicMock()
    db.query = mock_db_query

    for day in range(1, days + 1):
        # 1. Update user state based on the archetype's behavior
        behavior_fn(user, day, progress_history)

        # 2. Get Behavior Signal
        state_override = [user.competence, user.mood, user.fatigue, user.sleep_quality, user.stress, 0.0, 0.0]
        
        # Simulate getting historical features by replacing the DB query inside rl_engine
        import app.services.rl_engine as rl
        original_compute = rl._compute_historical_features
        
        def fake_compute(*args):
            return {
                "consecutive_hard_days":  sum(1 for r in progress_history[-3:] if r.actual_duration_min >= 60),
                "recent_completion_rate": 0.8,
                "avg_satisfaction_7d":    sum(r.satisfaction for r in progress_history[-7:] if r.satisfaction) / max(1, len(progress_history[-7:])),
                "fatigue_trend":          0.05 if user.fatigue > 0.7 else 0.0,
                "overwork_signal":        0.8 if user.fatigue > 0.8 else 0.0,
                "deep_work_streak":       sum(1 for r in progress_history[-3:] if r.action_label == "Deep Work"),
                "avg_daily_minutes":      60.0,
                "consecutive_same_action": 3 if len(progress_history) >= 3 and all(r.action_label == "Deep Work" for r in progress_history[-3:]) else 0,
                "last_action":            progress_history[-1].action_label if progress_history else "Normal",
                "ema_fatigue":            user.fatigue,
                "ema_mood":               user.mood,
                "burnout_warning":        user.fatigue > 0.8 and json.loads(user.rl_bandit_params).get("_recovery_debt", 0) > 10.0 if user.rl_bandit_params else False
            }
        rl._compute_historical_features = fake_compute
        
        try:
            signal = get_behavior_signal(user, state_override=state_override, db=db)
        finally:
            rl._compute_historical_features = original_compute
            
        action = signal["action_label"]
        params = _load_bandit_params(user)
        debt = params.get("_recovery_debt", 0.0)
        
        # 3. Apply consequence of the action (Mocking the user performing the task)
        # Deep Work generates high fatigue. Recovery reduces it.
        reward = 1.0
        duration = 60
        if action == "Deep Work":
            user.fatigue = min(1.0, user.fatigue + 0.15)
            user.mood = max(0.1, user.mood - 0.05)
            reward = 0.9 if user.fatigue < 0.8 else 0.4 # Reward drops if exhausted
        elif action in ("Recovery", "Light Review", "Rest"):
            user.fatigue = max(0.0, user.fatigue - 0.2)
            user.mood = min(1.0, user.mood + 0.1)
            duration = 30
            reward = 0.8
        else:
            user.fatigue = min(1.0, user.fatigue + 0.05)
            
        # 4. Save to history
        p = MagicMock()
        p.date = (datetime.date.today() - datetime.timedelta(days=days-day)).isoformat()
        p.completed = True
        p.action_label = action
        p.actual_duration_min = duration
        p.satisfaction = reward
        p.fatigue = user.fatigue
        p.mood = user.mood
        progress_history.append(p)
        if len(progress_history) > 14:
            progress_history.pop(0)

        # 5. Update Reward
        update_reward(db, user, action, reward)
        
        print(f"Day {day:02d} | Fatigue: {user.fatigue:.2f} | Mood: {user.mood:.2f} | Debt: {debt:5.2f} | Action: {action:12s} | Rew: {reward:.2f} | Warning: {signal.get('burnout_warning', False)}")

# --- Behaviors ---
def overachiever_behavior(user, day, hist):
    # Ignores fatigue, always tries to stay motivated. Doesn't sleep well if stressed.
    user.sleep_quality = 0.6 if user.fatigue > 0.6 else 0.9
    user.stress = 0.5
    if day == 1:
        user.fatigue = 0.1
        user.mood = 0.9

def erratic_behavior(user, day, hist):
    # High variance in sleep and mood
    user.sleep_quality = random.uniform(0.3, 1.0)
    user.mood += random.uniform(-0.3, 0.3)
    user.mood = max(0.1, min(1.0, user.mood))
    if day == 1:
        user.fatigue = 0.3

def normal_behavior(user, day, hist):
    # Standard user, sleeps okay, mood fluctuates mildly
    user.sleep_quality = 0.8
    if day % 7 in (6, 0): # Weekend
        user.fatigue = max(0.0, user.fatigue - 0.3)
        user.mood = min(1.0, user.mood + 0.2)
    if day == 1:
        user.fatigue = 0.2
        user.mood = 0.7

if __name__ == "__main__":
    simulate_archetype("The Overachiever", 30, overachiever_behavior)
    simulate_archetype("The Erratic", 30, erratic_behavior)
    simulate_archetype("The Normal", 30, normal_behavior)