# environment.py
# The actual RL environment logic for the Smart Study Session Optimizer
#
# The core idea: student has energy + focus, each action moves time forward by
# a fixed interval. Energy drains while studying, recovers on breaks.
# Focus tracks how well you can concentrate.
#
# NOTE: this works, don't touch it (the reward shaping took forever to tune)

import random
import math
from typing import Optional
from .models import Observation, StepResult

# Time constants (in minutes)
STUDY_INTERVAL = 5        # each "study" action = 5 minutes
SHORT_BREAK = 10          # short break duration
LONG_BREAK = 20           # long break duration

# Energy mechanics
ENERGY_DRAIN_BASE = 3.0       # energy lost per study interval
ENERGY_DRAIN_FOCUS_BONUS = 1.5  # extra drain when you're really focused
SHORT_BREAK_RECOVERY = 15.0
LONG_BREAK_RECOVERY = 35.0

# Focus mechanics - focus degrades the longer you study without a break
FOCUS_DECAY_RATE = 2.5        # focus lost per study interval
FOCUS_RECOVERY_SHORT = 20.0
FOCUS_RECOVERY_LONG = 45.0

# Thresholds
LOW_ENERGY_THRESHOLD = 40.0   # below this = bad idea to keep studying
HIGH_FOCUS_THRESHOLD = 60.0   # above this = productive


class StudyEnvironment:
    """
    The main environment class. Simulates a student's study session.

    State space:
        - energy (0-100): physical/mental stamina
        - focus (0-100): concentration level right now
        - time_studied: cumulative study minutes
        - session_minutes: total session time elapsed
        - consecutive_study_minutes: streak without a break
        - break_count: breaks taken
        - productive_minutes: minutes with focus > 60

    Action space:
        0 = keep studying (5 min)
        1 = short break (10 min)
        2 = long break (20 min)
    """

    @staticmethod
    def _strict_score(value: float) -> float:
        """Clamp a score strictly into [0.1, 0.9].
        The validator rejects 0.0 and 1.0, so we use a very safe margin."""
        return max(0.1, min(0.9, round(float(value), 4)))

    def __init__(self, task_id: str = "maximize_productivity", max_session_minutes: int = 480):
        self.task_id = task_id
        self.max_session_minutes = max_session_minutes

        # task-specific overrides
        # TODO: maybe make this a config dict later
        if task_id == "survive_8_hours":
            self.max_session_minutes = 480
        elif task_id == "maximize_productivity":
            self.max_session_minutes = 240
        elif task_id == "smart_break_taker":
            self.max_session_minutes = 400

        self._state = {}
        self.reset()

    def reset(self) -> Observation:
        """Start fresh. Called at the beginning of each episode."""
        print(f"[ENV] Resetting environment, task={self.task_id}")

        # small random variation so the agent has to generalize
        e = 85.0 + random.uniform(-10, 10)   # e = energy (yeah, short name, deal with it)
        f = 80.0 + random.uniform(-10, 10)   # f = focus

        self._state = {
            "energy": min(100.0, max(0.0, e)),
            "focus": min(100.0, max(0.0, f)),
            "time_studied": 0.0,
            "session_minutes": 0.0,
            "consecutive_study_minutes": 0.0,
            "break_count": 0,
            "productive_minutes": 0.0,
            # tracking stuff for graders
            "good_breaks": 0,   # breaks taken when energy < 40
            "total_breaks": 0,
            "done": False,
            "step_count": 0,
        }

        return self._get_obs()

    def step(self, action: int) -> StepResult:
        """
        Execute one action and return the result.
        action: 0=study, 1=short_break, 2=long_break
        """
        if self._state["done"]:
            # this shouldn't happen but just in case
            print("[ENV] WARNING: step() called after episode done")
            return StepResult(
                observation=self._get_obs(),
                reward=0.0,
                done=True,
                info=self._get_info()
            )

        s = self._state   # shorthand (this works, don't touch it)
        s["step_count"] += 1

        reward = 0.0

        if action == 0:  # study
            reward = self._do_study()
        elif action == 1:  # short break
            reward = self._do_short_break()
        elif action == 2:  # long break
            reward = self._do_long_break()
        else:
            print(f"[ENV] Unknown action {action}, defaulting to study")
            reward = self._do_study()

        # check termination conditions
        done = self._check_done()
        s["done"] = done

        if done:
            # add a terminal reward based on task
            reward += self._terminal_reward()

        obs = self._get_obs()
        info = self._get_info()

        print(f"[ENV] Step {s['step_count']}: action={action}, reward={reward:.3f}, "
              f"energy={s['energy']:.1f}, focus={s['focus']:.1f}, done={done}")

        return StepResult(observation=obs, reward=reward, done=done, info=info)

    def _do_study(self) -> float:
        """Handle the study action. Returns reward for this step."""
        s = self._state

        # energy drains faster when you're really focused (makes sense actually)
        focus_factor = s["focus"] / 100.0
        drain = ENERGY_DRAIN_BASE + (ENERGY_DRAIN_FOCUS_BONUS * focus_factor)
        s["energy"] = max(0.0, s["energy"] - drain)

        # focus decays the longer you study without a break
        consecutive_factor = min(s["consecutive_study_minutes"] / 60.0, 1.5)
        focus_drop = FOCUS_DECAY_RATE * (1 + consecutive_factor * 0.5)
        s["focus"] = max(0.0, s["focus"] - focus_drop)

        # advance time
        s["time_studied"] += STUDY_INTERVAL
        s["session_minutes"] += STUDY_INTERVAL
        s["consecutive_study_minutes"] += STUDY_INTERVAL

        # check if this was a productive interval
        if s["focus"] >= HIGH_FOCUS_THRESHOLD:
            s["productive_minutes"] += STUDY_INTERVAL

        # reward calculation
        reward = 0.0

        # positive reward for studying when you have energy and focus
        if s["energy"] > LOW_ENERGY_THRESHOLD and s["focus"] > HIGH_FOCUS_THRESHOLD:
            reward += 1.0
        elif s["energy"] > LOW_ENERGY_THRESHOLD:
            reward += 0.3
        else:
            # studying while exhausted - penalize slightly
            reward -= 0.5

        # small bonus just for getting through time (survival reward)
        reward += 0.1

        return reward

    def _do_short_break(self) -> float:
        """Short break: 10 minutes, moderate recovery."""
        s = self._state

        e_before = s["energy"]
        s["total_breaks"] += 1
        s["break_count"] += 1

        # check if this was a "smart" break (energy was low)
        if e_before < LOW_ENERGY_THRESHOLD:
            s["good_breaks"] += 1

        s["energy"] = min(100.0, s["energy"] + SHORT_BREAK_RECOVERY)
        s["focus"] = min(100.0, s["focus"] + FOCUS_RECOVERY_SHORT)
        s["consecutive_study_minutes"] = 0.0

        # time passes but no studying happens
        s["session_minutes"] += SHORT_BREAK

        # reward: good if you needed the break, meh if you didn't
        if e_before < LOW_ENERGY_THRESHOLD:
            reward = 0.8   # smart move
        elif e_before < 60.0:
            reward = 0.2
        else:
            reward = -0.3  # wasted time when you could have been studying

        return reward

    def _do_long_break(self) -> float:
        """Long break: 20 minutes, full recovery but costs more time."""
        s = self._state

        e_before = s["energy"]
        s["total_breaks"] += 1
        s["break_count"] += 1

        if e_before < LOW_ENERGY_THRESHOLD:
            s["good_breaks"] += 1

        s["energy"] = min(100.0, s["energy"] + LONG_BREAK_RECOVERY)
        s["focus"] = min(100.0, s["focus"] + FOCUS_RECOVERY_LONG)
        s["consecutive_study_minutes"] = 0.0
        s["session_minutes"] += LONG_BREAK

        # long break is only a good call when you're really drained
        if e_before < 25.0:
            reward = 1.0   # definitely needed this
        elif e_before < LOW_ENERGY_THRESHOLD:
            reward = 0.5
        else:
            reward = -0.6  # overkill, wasted too much time

        return reward

    def _check_done(self) -> bool:
        """Is the episode over?"""
        s = self._state

        # ran out of energy - bad
        if s["energy"] <= 0.0:
            print("[ENV] Episode ended: energy depleted!")
            return True

        # reached max session time
        if s["session_minutes"] >= self.max_session_minutes:
            print(f"[ENV] Episode ended: session complete ({s['session_minutes']} min)")
            return True

        # TODO: add a step limit as safety net for infinite loops

        return False

    def _terminal_reward(self) -> float:
        """Bonus/penalty given at the very end of the episode."""
        s = self._state
        terminal = 0.0

        if self.task_id == "survive_8_hours":
            # reward for surviving with energy left
            terminal = (s["energy"] / 100.0) * 5.0
        elif self.task_id == "maximize_productivity":
            # big reward for productive minutes
            terminal = (s["productive_minutes"] / 240.0) * 5.0
        elif self.task_id == "smart_break_taker":
            if s["total_breaks"] > 0:
                ratio = s["good_breaks"] / s["total_breaks"]
                terminal = ratio * 5.0

        return terminal

    def _get_obs(self) -> Observation:
        """Build the current observation from internal state."""
        s = self._state
        return Observation(
            energy=round(s["energy"], 2),
            focus=round(s["focus"], 2),
            time_studied=round(s["time_studied"], 2),
            session_minutes=round(s["session_minutes"], 2),
            consecutive_study_minutes=round(s["consecutive_study_minutes"], 2),
            break_count=s["break_count"],
            productive_minutes=round(s["productive_minutes"], 2),
        )

    def _get_info(self) -> dict:
        """Extra info for graders and debugging.
        CRITICAL: every score MUST be strictly in (0, 1) — never 0.0 or 1.0."""
        s = self._state
        total_breaks = max(s["total_breaks"], 1)

        # raw scores
        survive_score = s["energy"] / 100.0
        productivity_score = s["productive_minutes"] / 240.0

        if total_breaks > 0:
            break_score = s["good_breaks"] / total_breaks
        else:
            break_score = 0.0

        # strictly clamp into (0,1) — _strict_score guarantees (0.01, 0.99)
        scores = {
            "survive_8_hours": self._strict_score(survive_score),
            "maximize_productivity": self._strict_score(productivity_score),
            "smart_break_taker": self._strict_score(break_score),
        }

        current_score = scores.get(self.task_id, 0.5)

        # FINAL SAFETY NET: if anything is still 0.0 or 1.0, fix it
        for k in scores:
            if scores[k] <= 0.05:
                scores[k] = 0.1
            if scores[k] >= 0.95:
                scores[k] = 0.9
        if current_score <= 0.05:
            current_score = 0.1
        if current_score >= 0.95:
            current_score = 0.9

        return {
            "task_id": self.task_id,
            "score": current_score,
            "step_count": s["step_count"],
            "total_breaks": s["total_breaks"],
            "good_breaks": s["good_breaks"],
            "energy_remaining": round(s["energy"], 2),
            "productive_minutes": round(s["productive_minutes"], 2),
            "all_task_scores": scores,
        }

    def get_state(self) -> dict:
        """Return full internal state (for the /state endpoint)."""
        return {**self._state, "task_id": self.task_id}