"""Background worker for training runs."""

import argparse
import json
import signal
import sys
import traceback
from pathlib import Path

# Import tasks to register them
import vespaembed.tasks  # noqa: F401
from vespaembed.core.config import (
    DataConfig,
    LoraConfig,
    OutputConfig,
    TrainingConfig,
    TrainingHyperparameters,
    UnslothConfig,
)
from vespaembed.core.trainer import VespaEmbedTrainer
from vespaembed.db import update_run_status
from vespaembed.enums import RunStatus
from vespaembed.models.export import export_model
from vespaembed.utils.logging import logger


class TrainingWorker:
    """Worker that executes training in a subprocess."""

    def __init__(self, run_id: int, config: dict):
        self.run_id = run_id
        self.config = config
        self.trainer = None
        self.stopped = False

        # Set up signal handlers
        signal.signal(signal.SIGTERM, self._handle_signal)
        signal.signal(signal.SIGINT, self._handle_signal)

    def _handle_signal(self, signum, frame):
        """Handle termination signals."""
        logger.info(f"Received signal {signum}, stopping training...")
        self.stopped = True
        if self.trainer:
            # Trainer will check for stop flag
            pass
        sys.exit(0)

    def _send_update(self, update_type: str, data: dict):
        """Send update to the web server via file."""
        update_dir = Path.home() / ".vespaembed" / "updates"
        update_dir.mkdir(parents=True, exist_ok=True)

        update_file = update_dir / f"run_{self.run_id}.jsonl"
        update = {"type": update_type, "run_id": self.run_id, **data}

        with open(update_file, "a") as f:
            f.write(json.dumps(update) + "\n")

    def _log_config(self):
        """Log training configuration to UI."""
        c = self.config
        lines = [
            "=" * 50,
            "Training Configuration",
            "=" * 50,
            f"  Base Model:      {c.get('base_model', '--')}",
            f"  Task:            {c.get('task', '--')}",
        ]

        if c.get("loss_variant"):
            lines.append(f"  Loss Variant:    {c['loss_variant']}")

        # Data source
        if c.get("train_filename"):
            lines.append(f"  Training Data:   {c['train_filename'].split('/')[-1]}")
        elif c.get("hf_dataset"):
            lines.append(f"  Training Data:   {c['hf_dataset']}")

        # Training params
        lines.extend(
            [
                f"  Epochs:          {c.get('epochs', 3)}",
                f"  Batch Size:      {c.get('batch_size', 32)}",
                f"  Learning Rate:   {c.get('learning_rate', 2e-5)}",
                f"  Optimizer:       {c.get('optimizer', 'adamw_torch')}",
                f"  Scheduler:       {c.get('scheduler', 'linear')}",
                f"  Warmup Ratio:    {c.get('warmup_ratio', 0.1)}",
                f"  Weight Decay:    {c.get('weight_decay', 0.01)}",
            ]
        )

        # Precision
        if c.get("bf16"):
            lines.append("  Precision:       BF16")
        elif c.get("fp16"):
            lines.append("  Precision:       FP16")
        else:
            lines.append("  Precision:       FP32")

        # Optional features
        if c.get("gradient_checkpointing"):
            lines.append("  Grad Checkpoint: Enabled")
        if c.get("gradient_accumulation_steps", 1) > 1:
            lines.append(f"  Grad Accum:      {c['gradient_accumulation_steps']}")

        # LoRA
        if c.get("lora_enabled"):
            lines.append(f"  LoRA:            r={c.get('lora_r', 64)}, alpha={c.get('lora_alpha', 128)}")

        # Unsloth
        if c.get("unsloth_enabled"):
            lines.append(f"  Unsloth:         Enabled ({c.get('unsloth_save_method', 'merged_16bit')})")

        # Matryoshka
        if c.get("matryoshka_dims"):
            lines.append(f"  Matryoshka:      {c['matryoshka_dims']}")

        lines.append("=" * 50)

        # Send each line as a log message
        for line in lines:
            self._send_update("log", {"message": line})

    def run(self):
        """Execute the training run."""
        try:
            logger.info(f"Starting training run {self.run_id}")
            self._send_update("status", {"status": "running"})

            # Log configuration to UI
            self._log_config()

            # Determine data source - file upload or HuggingFace dataset
            train_source = self.config.get("train_filename") or self.config.get("hf_dataset")
            eval_source = self.config.get("eval_filename")
            eval_split = None

            # For HuggingFace datasets, eval can come from a different split
            hf_eval_split = self.config.get("hf_eval_split")
            if self.config.get("hf_dataset") and hf_eval_split:
                # Use the same dataset but different split for eval
                eval_source = self.config.get("hf_dataset")
                eval_split = hf_eval_split

            # Build data config
            data_config = DataConfig(
                train=train_source,
                eval=eval_source,
                eval_split_pct=self.config.get("eval_split_pct"),
                subset=self.config.get("hf_subset"),
                split=self.config.get("hf_train_split", "train"),
                eval_split=eval_split,
            )

            # Parse matryoshka_dims if present (comes as comma-separated string from UI)
            matryoshka_dims = None
            if self.config.get("matryoshka_dims"):
                dims_str = self.config["matryoshka_dims"]
                if isinstance(dims_str, str):
                    matryoshka_dims = [int(d.strip()) for d in dims_str.split(",") if d.strip()]
                elif isinstance(dims_str, list):
                    matryoshka_dims = dims_str

            # Parse lora_target_modules (comes as comma-separated string from UI)
            modules_str = self.config.get("lora_target_modules", "query, key, value, dense")
            if isinstance(modules_str, str):
                lora_target_modules = [m.strip() for m in modules_str.split(",") if m.strip()]
            else:
                lora_target_modules = modules_str  # Already a list

            # Build LoRA config
            lora_config = LoraConfig(
                enabled=self.config.get("lora_enabled", False),
                r=self.config.get("lora_r", 64),
                alpha=self.config.get("lora_alpha", 128),
                dropout=self.config.get("lora_dropout", 0.1),
                target_modules=lora_target_modules,
            )

            # Build Unsloth config
            unsloth_config = UnslothConfig(
                enabled=self.config.get("unsloth_enabled", False),
                save_method=self.config.get("unsloth_save_method", "merged_16bit"),
            )

            # Build training config with nested structure
            training_config = TrainingConfig(
                task=self.config["task"],
                base_model=self.config["base_model"],
                data=data_config,
                training=TrainingHyperparameters(
                    epochs=self.config.get("epochs", 3),
                    max_steps=self.config.get("max_steps"),
                    batch_size=self.config.get("batch_size", 32),
                    learning_rate=self.config.get("learning_rate", 2e-5),
                    warmup_ratio=self.config.get("warmup_ratio", 0.1),
                    weight_decay=self.config.get("weight_decay", 0.01),
                    fp16=self.config.get("fp16", False),
                    bf16=self.config.get("bf16", False),
                    eval_steps=self.config.get("eval_steps", 0.25),
                    save_steps=self.config.get("save_steps", 0.5),
                    logging_steps=self.config.get("logging_steps", 0.02),
                    gradient_accumulation_steps=self.config.get("gradient_accumulation_steps", 1),
                    optimizer=self.config.get("optimizer", "adamw_torch"),
                    scheduler=self.config.get("scheduler", "linear"),
                ),
                output=OutputConfig(
                    dir=self.config["output_dir"],  # Required - set by web app
                    push_to_hub=self.config.get("push_to_hub", False),
                    hf_username=self.config.get("hf_username"),
                ),
                lora=lora_config,
                unsloth=unsloth_config,
                max_seq_length=self.config.get("max_seq_length"),  # None = auto-detect
                gradient_checkpointing=self.config.get("gradient_checkpointing", False),
                matryoshka_dims=matryoshka_dims,
                loss_variant=self.config.get("loss_variant"),
            )

            # Create trainer with progress callback
            self.trainer = VespaEmbedTrainer(
                config=training_config,
                progress_callback=self._progress_callback,
            )

            # Run training
            self.trainer.train()

            # ONNX export (skip for Unsloth runs)
            if not self.config.get("unsloth_enabled", False):
                try:
                    self._send_update("log", {"message": "Exporting model to ONNX format..."})
                    final_model_path = Path(training_config.output.dir) / "final"
                    onnx_output = Path(training_config.output.dir) / "onnx"
                    export_model(str(final_model_path), str(onnx_output), format="onnx")
                    self._send_update("log", {"message": "ONNX export complete"})
                except Exception as e:
                    logger.warning(f"ONNX export failed: {e}")
                    self._send_update("log", {"message": f"ONNX export failed: {e} (training still succeeded)"})

            # Update status on completion
            update_run_status(self.run_id, RunStatus.COMPLETED)
            self._send_update("complete", {"output_dir": training_config.output.dir})
            logger.info(f"Training run {self.run_id} completed successfully")

        except Exception as e:
            logger.error(f"Training run {self.run_id} failed: {e}")
            logger.error(traceback.format_exc())
            update_run_status(self.run_id, RunStatus.ERROR, error_message=str(e))
            self._send_update("error", {"message": str(e)})
            sys.exit(1)

    def _progress_callback(self, progress: dict):
        """Callback for training progress updates."""
        progress_type = progress.get("type", "progress")

        if progress_type == "train_start":
            self._send_update("progress", progress)
            self._send_update(
                "log",
                {"message": f"Training started: {progress['total_steps']} steps, {progress['total_epochs']} epochs"},
            )
        elif progress_type == "train_end":
            self._send_update("progress", progress)
            elapsed = self._format_time(progress.get("elapsed_seconds", 0))
            self._send_update("log", {"message": f"Training completed in {elapsed}"})
        else:
            self._send_update("progress", progress)
            self._send_update("log", {"message": self._format_progress(progress)})

    def _format_time(self, seconds: float) -> str:
        """Format seconds into human readable time."""
        if seconds < 60:
            return f"{seconds:.0f}s"
        elif seconds < 3600:
            mins = int(seconds // 60)
            secs = int(seconds % 60)
            return f"{mins}m {secs}s"
        else:
            hours = int(seconds // 3600)
            mins = int((seconds % 3600) // 60)
            return f"{hours}h {mins}m"

    def _format_progress(self, progress: dict) -> str:
        """Format progress as a log message like tqdm."""
        step = progress.get("step", 0)
        total = progress.get("total_steps", 0)
        pct = progress.get("progress_pct", 0)
        loss = progress.get("loss")
        lr = progress.get("learning_rate")
        eta = progress.get("eta_seconds", 0)

        # Format: [Step 100/1000 (10%)] | Loss: 0.5234 | LR: 2.00e-05 | ETA: 5m 30s
        parts = [f"[Step {step}/{total} ({pct:.0f}%)]"]

        if loss is not None:
            parts.append(f"Loss: {loss:.4f}")

        # Add evaluation metrics if present
        eval_metrics = progress.get("eval_metrics", {})
        if eval_metrics:
            for key, value in eval_metrics.items():
                # Remove eval_ prefix and format nicely
                metric_name = key.replace("eval_", "").replace("_", " ").title()
                if isinstance(value, float):
                    parts.append(f"{metric_name}: {value:.4f}")
                else:
                    parts.append(f"{metric_name}: {value}")

        if lr is not None:
            parts.append(f"LR: {lr:.2e}")

        if eta > 0:
            parts.append(f"ETA: {self._format_time(eta)}")

        return " | ".join(parts)


def main():
    """Main entry point for the worker."""
    parser = argparse.ArgumentParser(description="VespaEmbed Training Worker")
    parser.add_argument("--run-id", type=int, required=True, help="Run ID")
    parser.add_argument("--config", type=str, required=True, help="Training config JSON")

    args = parser.parse_args()

    config = json.loads(args.config)
    worker = TrainingWorker(run_id=args.run_id, config=config)
    worker.run()


if __name__ == "__main__":
    main()
