"""RefreshWheel — main class for continual LLM fine-tuning."""
from __future__ import annotations

import json
import shutil
import tempfile
import uuid
from pathlib import Path
from typing import Any, Optional

from llm_refresh.buffer import BaseBuffer, create_buffer
from llm_refresh.config import AppConfig, get_config
from llm_refresh.metrics import ForgettingTracker
from llm_refresh.models import BufferStrategy, EvalResult, RefreshResult

try:
    import torch
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False

try:
    from transformers import AutoModelForCausalLM, AutoTokenizer
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False

try:
    from peft import LoraConfig, TaskType, get_peft_model
    HAS_PEFT = True
except ImportError:
    HAS_PEFT = False

try:
    from trl import SFTConfig, SFTTrainer
    HAS_TRL = True
except ImportError:
    HAS_TRL = False


class RefreshWheel:
    """
    Smart continual fine-tuning wrapper with replay buffer + forgetting metrics.

    Usage:
        rw = RefreshWheel(model_name="microsoft/phi-2")
        rw.set_anchor(anchor_examples)
        rw.add_data(new_examples)
        result = rw.refresh(epochs=1)
        print(rw.status())
    """

    def __init__(
        self,
        model_name: Optional[str] = None,
        buffer_strategy: BufferStrategy | str = BufferStrategy.RESERVOIR,
        config: Optional[AppConfig] = None,
        state_path: Optional[str | Path] = None,
    ):
        self.config = config or get_config()
        self._model_name = model_name or self.config.model.name
        self._state_path = Path(state_path) if state_path else None

        buf_cfg = self.config.buffer
        self.buffer: BaseBuffer = create_buffer(
            buffer_strategy,
            max_size=buf_cfg.max_size,
        )
        self.tracker = ForgettingTracker(krs_lambda=self.config.eval.krs_lambda)

        self._model: Any = None
        self._tokenizer: Any = None
        self._anchor_data: list[dict[str, Any]] = []
        self._new_data: list[dict[str, Any]] = []
        self._last_refresh: Optional[str] = None

    # ------------------------------------------------------------------
    # Model loading (lazy)
    # ------------------------------------------------------------------

    def load_model(self) -> None:
        """Load base model + LoRA adapter. Requires torch + transformers + peft."""
        if not HAS_TRANSFORMERS:
            raise RuntimeError("Install: pip install llm-refresh-wheel[transformers]")
        if not HAS_PEFT:
            raise RuntimeError("Install: pip install llm-refresh-wheel[peft]")

        from llm_refresh.utils import print_info
        print_info(f"Loading model: {self._model_name}")

        import torch
        dtype = torch.float16 if torch.cuda.is_available() else torch.float32

        self._tokenizer = AutoTokenizer.from_pretrained(self._model_name)
        if self._tokenizer.pad_token is None:
            self._tokenizer.pad_token = self._tokenizer.eos_token

        base = AutoModelForCausalLM.from_pretrained(
            self._model_name, torch_dtype=dtype, device_map="auto"
        )

        lora_cfg = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=self.config.model.rank,
            lora_alpha=self.config.model.lora_alpha,
            lora_dropout=self.config.model.lora_dropout,
            target_modules=self.config.model.target_modules,
        )
        self._model = get_peft_model(base, lora_cfg)
        print_info("Model loaded with LoRA adapter.")

    # ------------------------------------------------------------------
    # Data management
    # ------------------------------------------------------------------

    def add_data(self, examples: list[dict[str, Any]]) -> int:
        """Stage new examples for next refresh. Returns count added."""
        self._new_data.extend(examples)
        return len(examples)

    def set_anchor(self, examples: list[dict[str, Any]]) -> int:
        """Set the anchor evaluation set. Returns count set."""
        self._anchor_data = list(examples)
        return len(examples)

    # ------------------------------------------------------------------
    # Evaluation
    # ------------------------------------------------------------------

    def eval(self, dataset: Optional[list[dict[str, Any]]] = None) -> EvalResult:
        """
        Compute perplexity on `dataset` (defaults to anchor set).
        Requires a loaded model.
        """
        if self._model is None:
            self.load_model()
        from llm_refresh.hooks import compute_perplexity
        data = dataset if dataset is not None else self._anchor_data
        if not data:
            raise ValueError("No dataset provided and anchor set is empty.")
        return compute_perplexity(
            self._model,
            self._tokenizer,
            data,
            batch_size=self.config.eval.batch_size,
        )

    # ------------------------------------------------------------------
    # Refresh
    # ------------------------------------------------------------------

    def refresh(
        self,
        epochs: int = 1,
        dry_run: bool = False,
    ) -> RefreshResult:
        """
        Run one continual learning cycle:
        1. Eval on anchor → pre_eval
        2. Build training set: new data + replay samples
        3. Save temp checkpoint
        4. Train with SFTTrainer (if not dry_run)
        5. Eval on anchor → post_eval
        6. Compute BWT; rollback if forgetting detected
        7. Update buffer with new data
        8. Record and return RefreshResult
        """
        from llm_refresh.utils import print_info, print_warning, print_success

        if not self._anchor_data:
            raise ValueError("Set an anchor dataset first: rw.set_anchor(examples)")

        new_examples = list(self._new_data)
        self._new_data = []

        # 1. Pre-eval
        pre_eval: Optional[EvalResult] = None
        if not dry_run:
            if self._model is None:
                self.load_model()
            print_info("Running pre-training eval on anchor set...")
            pre_eval = self.eval(self._anchor_data)
            self.tracker.record_pre_eval(pre_eval)
        else:
            print_info("[dry-run] Skipping model load + eval.")

        # 2. Build training set
        replay_ratio = self.config.buffer.min_replay_ratio
        n_new = len(new_examples)
        n_replay = max(1, int(n_new * replay_ratio / (1 - replay_ratio))) if n_new else 0
        replay_samples = self.buffer.sample(n_replay) if not dry_run else []
        training_set = new_examples + replay_samples

        print_info(f"Training set: {len(new_examples)} new + {len(replay_samples)} replay = {len(training_set)} total")

        # 3. Save temp checkpoint
        checkpoint_dir: Optional[str] = None
        if not dry_run and self._model is not None and HAS_PEFT:
            checkpoint_dir = tempfile.mkdtemp(prefix="llm_refresh_ckpt_")
            self._model.save_pretrained(checkpoint_dir)
            print_info(f"Saved temp checkpoint → {checkpoint_dir}")

        # 4. Train
        rolled_back = False
        rollback_reason = ""
        post_eval: Optional[EvalResult] = None

        if not dry_run and training_set:
            if not HAS_TRL:
                raise RuntimeError("Install: pip install llm-refresh-wheel[train]")

            from datasets import Dataset as HFDataset

            train_dataset = HFDataset.from_list(training_set)
            t_cfg = self.config.training

            sft_args = SFTConfig(
                output_dir=checkpoint_dir or tempfile.mkdtemp(),
                num_train_epochs=epochs,
                per_device_train_batch_size=t_cfg.batch_size,
                gradient_accumulation_steps=t_cfg.gradient_accumulation_steps,
                learning_rate=t_cfg.learning_rate,
                warmup_ratio=t_cfg.warmup_ratio,
                max_seq_length=t_cfg.max_seq_length,
                logging_steps=10,
                save_strategy="no",
                report_to="none",
            )

            trainer = SFTTrainer(
                model=self._model,
                tokenizer=self._tokenizer,
                args=sft_args,
                train_dataset=train_dataset,
                dataset_text_field="text",
            )
            trainer.train()

        # 5. Post-eval
        if not dry_run:
            print_info("Running post-training eval on anchor set...")
            post_eval = self.eval(self._anchor_data)
            self.tracker.record_post_eval(post_eval)

        # 6. Compute BWT
        bwt = 0.0
        fwt = 0.0
        if pre_eval and post_eval:
            bwt = pre_eval.perplexity - post_eval.perplexity

        # Compute FWT from previous cycle
        history = self.tracker.history
        if len(history) >= 1 and pre_eval is not None:
            prev = history[-1]
            if prev.pre_eval:
                fwt = prev.pre_eval.perplexity - pre_eval.perplexity

        krs = self.tracker.compute_krs()

        # 7. Rollback if forgetting
        if (
            not dry_run
            and bwt < self.config.eval.forgetting_threshold
            and self.config.eval.auto_rollback
            and checkpoint_dir
            and self._model is not None
        ):
            print_warning(f"Forgetting detected (BWT={bwt:.4f}). Rolling back...")
            try:
                from peft import PeftModel
                self._model.load_adapter(checkpoint_dir, adapter_name="default")
                rolled_back = True
                rollback_reason = f"BWT {bwt:.4f} < threshold {self.config.eval.forgetting_threshold}"
                print_warning("Rollback complete.")
            except Exception as e:
                rollback_reason = f"Rollback failed: {e}"

        # Cleanup checkpoint
        if checkpoint_dir:
            try:
                shutil.rmtree(checkpoint_dir, ignore_errors=True)
            except Exception:
                pass

        # 8. Update buffer
        self.buffer.add(new_examples)

        # 9. Record result
        result = RefreshResult.new(
            examples_trained=len(training_set),
            replay_examples=len(replay_samples),
            new_examples=n_new,
            pre_eval=pre_eval,
            post_eval=post_eval,
            bwt=bwt,
            fwt=fwt,
            krs=krs,
            rolled_back=rolled_back,
            rollback_reason=rollback_reason,
        )
        self.tracker.record_result(result)
        self._last_refresh = result.timestamp

        if not dry_run:
            status_str = "rolled back" if rolled_back else "complete"
            print_success(f"Refresh {status_str}. BWT={bwt:.4f}, KRS={krs:.1f}")
        else:
            print_success(f"[dry-run] Refresh flow complete. Buffer size: {len(self.buffer)}")

        if self._state_path:
            self.save(self._state_path)

        return result

    # ------------------------------------------------------------------
    # Status
    # ------------------------------------------------------------------

    def status(self) -> dict[str, Any]:
        return {
            "model": self._model_name,
            "buffer": self.buffer.stats(),
            "metrics": self.tracker.summary(),
            "last_refresh": self._last_refresh,
            "anchor_size": len(self._anchor_data),
            "pending_examples": len(self._new_data),
        }

    # ------------------------------------------------------------------
    # Persistence (buffer + metrics, NOT model weights)
    # ------------------------------------------------------------------

    def save(self, path: Path | str) -> None:
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        # Save buffer
        self.buffer.save(path / "buffer.json")
        # Save metrics
        self.tracker.save(path / "metrics.json")
        # Save metadata
        meta = {
            "model_name": self._model_name,
            "last_refresh": self._last_refresh,
            "anchor_size": len(self._anchor_data),
        }
        (path / "meta.json").write_text(json.dumps(meta, indent=2))

    def load(self, path: Path | str) -> None:
        path = Path(path)
        buf_path = path / "buffer.json"
        if buf_path.exists():
            self.buffer.load(buf_path)
        met_path = path / "metrics.json"
        if met_path.exists():
            self.tracker = ForgettingTracker.load(met_path)
        meta_path = path / "meta.json"
        if meta_path.exists():
            meta = json.loads(meta_path.read_text())
            self._model_name = meta.get("model_name", self._model_name)
            self._last_refresh = meta.get("last_refresh")
