from dataclasses import dataclass
from typing import Callable, Sequence, TypeAlias

import numpy as np
from numpy.typing import NDArray

from adaptive_harmony import (
    JobNotifier,
    Logger,
    StageNotifier,
    StringThread,
    TokenizedThread,
    TrainingModel,
)
from adaptive_harmony.common import RecipeCallback
from adaptive_harmony.common.env_grpo import ENVGRPO
from adaptive_harmony.core.utils import (
    async_map,
    hash_hyperparams,
    log_args,
)
from adaptive_harmony.environment import EnvironmentFactory, TrajectoryScore
from adaptive_harmony.metric_logger import StdoutLogger

FloatArray: TypeAlias = NDArray[np.float32]


@dataclass
class Sample:
    sample: TokenizedThread
    logprobs: list[float]
    ref_logprobs: list[float]
    advantage: float
    kl_div: list[float]
    # for logging
    score: float
    gen_len: float


ENVGSPO_HYPERPARAMS = {
    "max_num_grpo_steps",
    "completions_per_sample",
    "lr",
    "lr_scheduler",
    "samples_per_batch",
    "samples_per_mini_batch",
    "mini_epochs_per_batch",
    "max_grad_norm",
    "clip_range",
    "kl_beta",
    "weight_decays",
    "skip_nan_gradients",
}


class ENVGSPO(ENVGRPO):
    @log_args
    @hash_hyperparams(include=ENVGSPO_HYPERPARAMS)
    def __init__(
        self,
        dataset: list[StringThread],
        model: TrainingModel,
        environment_factory: EnvironmentFactory,
        logger: Logger = StdoutLogger(),
        stage_notifier: StageNotifier = JobNotifier().stage_notifier("ENVGSPO Training"),
        callbacks: Sequence[RecipeCallback] = [],
        validation_dataset: list[StringThread] | None = None,
        validation_frequency: float = 0.2,
        max_num_grpo_steps: int | None = None,
        completions_per_sample=8,
        lr: float = 7.5e-7,
        lr_scheduler: Callable[[float], float] | None = None,
        samples_per_batch=128,
        samples_per_mini_batch=128,
        mini_epochs_per_batch=1,
        max_grad_norm=1.0,
        clip_range=0.01,
        kl_beta=0.1,
        weight_decays: float = 0.0,
        skip_nan_gradients: bool = False,
        restart_from_checkpoint: str | None = None,
        checkpoint_frequency: float = 0.2,
        data_seed: int = 42,
    ):
        super().__init__(
            dataset,
            model,
            environment_factory,
            logger,
            stage_notifier,
            callbacks,
            max_num_grpo_steps=max_num_grpo_steps,
            completions_per_sample=completions_per_sample,
            lr=lr,
            lr_scheduler=lr_scheduler,
            samples_per_batch=samples_per_batch,
            samples_per_mini_batch=samples_per_mini_batch,
            mini_epochs_per_batch=mini_epochs_per_batch,
            max_grad_norm=max_grad_norm,
            clip_range=clip_range,
            kl_beta=kl_beta,
            weight_decays=weight_decays,
            data_seed=data_seed,
        )

    async def gen_data(self, sample: StringThread) -> list[Sample]:
        # need to override gen data due to the single reward check
        async def generate_trajectory(
            prompt: StringThread,
        ) -> tuple[TokenizedThread, TrajectoryScore, int]:
            # this create the environment for the first turn.
            environment = self.environment_factory.create_environment(prompt.metadata)
            prompt = await environment.bootstrap_prompt(prompt)

            # Count assistant turns in the context (before generation)
            nb_context_assistant_turns = sum(1 for turn in prompt.get_turns() if turn.role == "assistant")

            string_trajectory = await self.model.generate(prompt)  # generate the first response from the agent.
            num_generated_turns = 1
            # we loop until the environment returns a score.
            # notice how the environment can return a score or a tool or user response.
            while not isinstance(
                environment_response := await environment.react_to(string_trajectory),
                TrajectoryScore,
            ):
                for env_role, env_content in environment_response:
                    if not isinstance(env_content, str):
                        raise ValueError(f"env_content should be a str, got {env_content}")
                    if env_role == "user":
                        string_trajectory = string_trajectory.user(env_content)
                    elif env_role == "tool":
                        string_trajectory = string_trajectory.tool(env_content)
                    else:
                        raise ValueError
                string_trajectory = await self.model.generate(string_trajectory)
                num_generated_turns += 1

            tokenized_trajectory = (
                await self.model.tokenize_thread(string_trajectory)
            ).with_weight_assistant_turns_from_index(nb_context_assistant_turns)

            return tokenized_trajectory, environment_response, num_generated_turns

        assert self.model_ref is not None, "Calling `gen_data` before reference model has been set"

        trajs_and_scores = await async_map(generate_trajectory, [sample] * self.completions_per_sample)
        all_samples = [traj for traj, _, _ in trajs_and_scores]
        logprobs = await async_map(self.model.logprobs_per_token, all_samples)
        ref_logprobs = await async_map(self.model_ref.logprobs_per_token, all_samples)

        all_trajectory_scores = [score for _, score, _ in trajs_and_scores]
        assert all(len(traj_score.scores) == 1 for traj_score in all_trajectory_scores)
        all_scores = np.array(
            [traj_score.scores[0].score for traj_score in all_trajectory_scores],
            dtype=np.float32,
        )

        advantages: FloatArray = all_scores - all_scores.mean()
        advantages /= advantages.std() + 1e-8

        kl = [
            (np.array(lp, dtype=np.float32) - np.array(ref_lp, dtype=np.float32)).tolist()
            for lp, ref_lp in zip(logprobs, ref_logprobs)
        ]

        samples = []
        for i in range(len(logprobs)):
            samples.append(
                Sample(
                    sample=all_samples[i],
                    logprobs=logprobs[i],
                    ref_logprobs=ref_logprobs[i],
                    advantage=advantages[i],
                    kl_div=kl[i],
                    score=all_trajectory_scores[i].cumulative_score,
                    gen_len=all_samples[i].len_last_turn(),
                )
            )
        return samples

    async def train_sample(self, sample: Sample):
        await self.model.train_gspo(
            sample.sample,
            sample.logprobs,
            sample.ref_logprobs,
            advantage=[sample.advantage],
            left_clip=self.clip_range,
            right_clip=self.clip_range,
            kl_beta=self.kl_beta,
        )
