"""Independent H2 GPU learner, based on 4fa48dd combined-action residual TD3."""
from __future__ import annotations

import copy
import json
import os
from pathlib import Path
import queue
import random
import signal
import time
import traceback

import numpy as np
import torch
from torch import nn

from .action import ResidualActionAdapter
from .config import RLConfig
from .data import ACTOR_INPUT_DIM, OBS_DIM, FORMAT, REWARD_PROTOCOL, ActionContract, EpisodeBatch
from .models import Actor, Critic, RunningNormalizer, flatten_parameters
from .replay import ReplayBuffer


def load_checkpoint(path, contract: ActionContract | None = None):
    # Checkpoints are local trusted training artifacts, never received over ZMQ.
    saved = torch.load(Path(path), map_location="cpu", weights_only=False)
    if not isinstance(saved, dict) or saved.get("format") != FORMAT:
        raise ValueError("requires H2 real residual v3 checkpoint; old H2/G1 checkpoints are unsupported")
    if saved.get("reward_protocol") != REWARD_PROTOCOL:
        raise ValueError("checkpoint reward protocol mismatch")
    if contract is not None:
        contract.validate_saved(saved.get("contract"))
    RLConfig(**saved["config"])
    return saved


class TD3Learner:
    def __init__(self, config: RLConfig, contract: ActionContract, *, device=None):
        self.config, self.contract = config, contract
        # CPU is a dependency-injection seam for tests; the process entrypoint
        # always supplies config.learner_device and never falls back from CUDA.
        self.device = torch.device(device or config.learner_device)
        if self.device.type == "cuda" and not torch.cuda.is_available():
            raise RuntimeError("H2 residual learner requires CUDA")
        random.seed(config.seed)
        np.random.seed(config.seed)
        torch.manual_seed(config.seed)
        self.actor = Actor().to(self.device)
        self.actor_target = copy.deepcopy(self.actor).eval()
        self.critics = [Critic().to(self.device), Critic().to(self.device)]
        self.critic_targets = [copy.deepcopy(c).eval() for c in self.critics]
        options = dict(betas=config.optimizer_betas, eps=config.optimizer_eps, weight_decay=config.weight_decay)
        self.actor_optimizer = torch.optim.AdamW(self.actor.parameters(), lr=config.actor_learning_rate, **options)
        self.critic_optimizers = [torch.optim.AdamW(c.parameters(), lr=config.critic_learning_rate, **options)
                                  for c in self.critics]
        self.normalizer = RunningNormalizer(ACTOR_INPUT_DIM, min_std=config.normalizer_min_std,
                                            clip=config.normalizer_clip)
        self.adapter = ResidualActionAdapter(contract)
        self.replay = ReplayBuffer(config.replay_capacity, config.seed)
        self.balls = self.warmup_balls = self.online_balls = self.transitions = 0
        self.critic_updates = self.actor_updates = self.warmup_updates = self.online_updates = 0
        self.pending_updates = self.policy_version = 0
        self.collector_rng_state = None
        self.last_summary = {}
        self.metrics = {}

    @property
    def warmup_complete(self):
        return (self.warmup_balls == self.config.warmup_balls and
                self.warmup_updates >= self.config.critic_warmup_steps and self.normalizer.frozen)

    def ingest(self, episode: EpisodeBatch):
        episode.validate()
        if episode.sequence != self.balls + 1:
            raise ValueError("episode sequence gap; refusing partial/misordered training")
        expected = "warmup" if self.warmup_balls < self.config.warmup_balls else "online"
        if episode.phase != expected or (expected == "online" and not self.warmup_complete):
            raise ValueError("episode training phase mismatch")
        if self.online_balls >= self.config.online_balls:
            raise ValueError("training ball budget exceeded")
        if expected == "warmup":
            for item in episode.transitions:
                self.normalizer.update(np.concatenate((item.base_obs, item.base_action)))
            self.warmup_balls += 1
            if self.warmup_balls == self.config.warmup_balls:
                self.normalizer.freeze()
        else:
            self.online_balls += 1
            self.pending_updates += len(episode.transitions) * self.config.critic_updates_per_step
        self.replay.add_episode(episode, self.config.n_step, self.config.discount)
        self.balls += 1
        self.transitions += len(episode.transitions)
        self.collector_rng_state = episode.collector_rng_state
        self.last_summary = dict(episode.summary, sent_to_learner=True)

    def _soft_update(self, source, target):
        with torch.no_grad():
            for src, dst in zip(source.parameters(), target.parameters()):
                dst.lerp_(src, self.config.target_tau)

    def update_once(self, *, critic_only=False):
        if self.replay.size < self.config.batch_size:
            raise RuntimeError("valid warmup balls did not produce a complete replay batch")
        c = self.config
        batch = {k: torch.as_tensor(v, device=self.device) for k, v in self.replay.sample(c.batch_size).items()}
        norm = self.normalizer.normalize_tensor
        inputs = norm(torch.cat((batch["base_obs"], batch["combined_action"]), dim=1))
        with torch.no_grad():
            next_inputs = norm(torch.cat((batch["next_base_obs"], batch["next_base_action"]), dim=1))
            residual = self.actor_target(next_inputs)
            if not critic_only:
                noise = (torch.randn_like(residual) * c.target_noise).clamp(-c.target_noise_clip, c.target_noise_clip)
                residual = (residual + noise).clamp(-1, 1)
            combined = self.adapter.combine_tensor(batch["next_base_action"], residual)
            target_inputs = norm(torch.cat((batch["next_base_obs"], combined), dim=1))
            q = torch.minimum(*(target(target_inputs) for target in self.critic_targets))
            target = batch["reward_n"] + batch["discount_n"] * batch["bootstrap_mask"] * q
        losses = []
        for critic, optimizer in zip(self.critics, self.critic_optimizers):
            loss = nn.functional.mse_loss(critic(inputs), target)
            if not torch.isfinite(loss):
                raise RuntimeError("nonfinite critic loss")
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            nn.utils.clip_grad_norm_(critic.parameters(), c.gradient_clip, error_if_nonfinite=True)
            optimizer.step()
            losses.append(float(loss.detach().cpu()))
        self.critic_updates += 1
        if critic_only:
            self.warmup_updates += 1
            for critic, target_net in zip(self.critics, self.critic_targets):
                self._soft_update(critic, target_net)
        else:
            self.online_updates += 1
            self.pending_updates -= 1
            if self.online_updates % c.actor_update_delay == 0:
                for critic in self.critics:
                    critic.requires_grad_(False)
                actor_inputs = torch.cat((batch["base_obs"], batch["base_action"]), dim=1)
                residual = self.actor(norm(actor_inputs))
                action = self.adapter.combine_tensor(actor_inputs[:, OBS_DIM:], residual)
                q_loss = -self.critics[0](norm(torch.cat((batch["base_obs"], action), dim=1))).mean()
                regularization = c.actor_residual_l2 * residual.square().mean()
                actor_loss = q_loss + regularization
                if not torch.isfinite(actor_loss):
                    raise RuntimeError("nonfinite actor loss")
                self.actor_optimizer.zero_grad(set_to_none=True)
                actor_loss.backward()
                nn.utils.clip_grad_norm_(self.actor.parameters(), c.gradient_clip, error_if_nonfinite=True)
                self.actor_optimizer.step()
                for critic in self.critics:
                    critic.requires_grad_(True)
                self._soft_update(self.actor, self.actor_target)
                for critic, target_net in zip(self.critics, self.critic_targets):
                    self._soft_update(critic, target_net)
                self.actor_updates += 1
                self.metrics["actor_loss"] = float(actor_loss.detach().cpu())
                self.metrics["actor_q_loss"] = float(q_loss.detach().cpu())
                self.metrics["actor_regularization"] = float(regularization.detach().cpu())
        self.metrics.update(critic_1_loss=losses[0], critic_2_loss=losses[1])

    def publish(self, sync):
        self.policy_version = sync.publish(flatten_parameters(self.actor), self.normalizer.mean,
                                           self.normalizer.std)

    def state_dict(self):
        counters = {name: getattr(self, name) for name in (
            "balls", "warmup_balls", "online_balls", "transitions", "critic_updates", "actor_updates",
            "warmup_updates", "online_updates", "pending_updates", "policy_version")}
        return dict(format=FORMAT, reward_protocol=REWARD_PROTOCOL,
                    source_commit="4fa48dd2672bd52d9349fb3f7902755d665aa725",
                    config=self.config.to_dict(), contract=self.contract.to_dict(), counters=counters,
                    actor=self.actor.state_dict(), actor_target=self.actor_target.state_dict(),
                    critics=[c.state_dict() for c in self.critics],
                    critic_targets=[c.state_dict() for c in self.critic_targets],
                    actor_optimizer=self.actor_optimizer.state_dict(),
                    critic_optimizers=[o.state_dict() for o in self.critic_optimizers],
                    normalizer=self.normalizer.state_dict(), replay=self.replay.state_dict(),
                    numpy_rng=np.random.get_state(), python_rng=random.getstate(), torch_rng=torch.get_rng_state(),
                    cuda_rng=torch.cuda.get_rng_state_all() if self.device.type == "cuda" else [],
                    collector_rng=self.collector_rng_state, last_episode=self.last_summary,
                    metrics=self.metrics)

    def save(self, path):
        path = Path(path)
        temporary = path.with_suffix(path.suffix + ".tmp")
        with temporary.open("wb") as stream:
            torch.save(self.state_dict(), stream)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)

    def restore(self, saved):
        self.contract.validate_saved(saved["contract"])
        if RLConfig(**saved["config"]) != self.config:
            raise ValueError("resume requires matching resolved H2 training config")
        self.actor.load_state_dict(saved["actor"])
        self.actor_target.load_state_dict(saved["actor_target"])
        for group, key in ((self.critics, "critics"), (self.critic_targets, "critic_targets"),
                           (self.critic_optimizers, "critic_optimizers")):
            for item, state in zip(group, saved[key]):
                item.load_state_dict(state)
        self.actor_optimizer.load_state_dict(saved["actor_optimizer"])
        norm = saved["normalizer"]
        self.normalizer.mean = np.asarray(norm["mean"], dtype=np.float64)
        self.normalizer.m2 = np.asarray(norm["m2"], dtype=np.float64)
        self.normalizer.count, self.normalizer.frozen = int(norm["count"]), bool(norm["frozen"])
        self.replay.load_state_dict(saved["replay"])
        for name, value in saved["counters"].items():
            setattr(self, name, int(value))
        np.random.set_state(saved["numpy_rng"])
        random.setstate(saved["python_rng"])
        torch.set_rng_state(saved["torch_rng"])
        if self.device.type == "cuda" and saved["cuda_rng"]:
            torch.cuda.set_rng_state_all(saved["cuda_rng"])
        self.collector_rng_state = saved["collector_rng"]
        self.last_summary, self.metrics = saved["last_episode"], saved["metrics"]


def learner_main(config, contract, episodes, sync, status, stop_event, run_dir, resume_path=None):
    """Sequential ingestion makes every checkpoint an exact ball/update boundary."""
    root = Path(run_dir)
    learner = None
    # The parent coordinates terminal Ctrl+C and drains every accepted episode.
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    try:
        torch.set_num_threads(1)
        learner = TD3Learner(config, contract)
        if resume_path:
            learner.restore(load_checkpoint(resume_path, contract))
            sync.version.value = learner.policy_version
        learner.publish(sync)
        last_metrics_s = 0.0

        def progress():
            nonlocal last_metrics_s
            status.ingested_balls.value = learner.balls
            status.transitions.value = learner.transitions
            status.critic_updates.value = learner.critic_updates
            status.actor_updates.value = learner.actor_updates
            status.pending_updates.value = learner.pending_updates + (
                max(0, config.critic_warmup_steps - learner.warmup_updates)
                if learner.warmup_balls == config.warmup_balls else 0)
            status.warmup_done.value = int(learner.warmup_complete)
            now = time.monotonic()
            if now - last_metrics_s >= 1:
                with (root / "learner.jsonl").open("a") as log:
                    log.write(json.dumps(dict(monotonic_s=now, **status.snapshot(), **learner.metrics)) + "\n")
                last_metrics_s = now

        # Publish restored counters and warmup state before admitting the first ball.
        progress()
        status.ready.value = 1
        while True:
            # Always finish the current ball's debt before consuming another;
            # other processes keep collecting with the last complete actor.
            if learner.warmup_balls == config.warmup_balls and not learner.warmup_complete:
                status.stage.value = 2
                while not learner.warmup_complete:
                    started = time.monotonic()
                    learner.update_once(critic_only=True)
                    status.last_update_s.value = time.monotonic() - started
                    progress()
                learner.publish(sync)
            if (learner.warmup_complete and learner.online_balls == 0 and
                    not (root / "checkpoint_warmup.pt").exists()):
                status.stage.value = 4
                learner.save(root / "checkpoint_warmup.pt")
            status.stage.value = 3
            while learner.pending_updates:
                started = time.monotonic()
                previous_actor_updates = learner.actor_updates
                learner.update_once()
                status.last_update_s.value = time.monotonic() - started
                if learner.actor_updates != previous_actor_updates:
                    learner.publish(sync)
                progress()
            progress()
            if learner.online_balls and learner.online_balls % config.checkpoint_interval_balls == 0:
                path = root / f"checkpoint_online_ball_{learner.online_balls:06d}.pt"
                if not path.exists():
                    status.stage.value = 4
                    learner.save(path)
            if learner.online_balls == config.online_balls:
                break
            try:
                status.stage.value = 1
                episode = episodes.get(timeout=0.1)
            except queue.Empty:
                if stop_event.is_set() and learner.balls >= status.stop_after_balls.value:
                    break
                continue
            learner.ingest(episode)
            progress()
        status.stage.value = 4
        learner.save(root / "checkpoint_final.pt")
        status.stage.value = 5
        status.done.value = 1
    except BaseException:
        status.failed.value = 1
        status.stage.value = 6
        (root / "learner_error.txt").write_text(traceback.format_exc())
        if learner is not None:
            try:
                learner.save(root / "checkpoint_interrupted.pt")
            except Exception:
                pass
    finally:
        status.ready.value = 0