"""远程部署/导出入口 — 在算力节点的训练容器里执行。"""
import json
import os
import shutil
from pathlib import Path

os.environ["PYTORCH_NO_FLASH"] = "1"
os.environ["MACA_MPS_MODE"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"

_DATA_DIR = Path(os.environ.get("COMPUTE_NODE_REMOTE_DATA_DIR", "/root/Fine-tuning/backend/data"))
_ADAPTERS_DIR = _DATA_DIR / "adapters"


async def run_remote_export(job_id: str, merge_with_base: bool = False, export_format: str = "safetensors") -> dict:
    """在远程容器里合并 adapter 或导出 GGUF。"""
    adapter_path = _ADAPTERS_DIR / job_id
    if not adapter_path.exists():
        return {"error": f"Adapter not found: {adapter_path}"}

    output_path = _ADAPTERS_DIR / f"{job_id}_merged"

    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        if merge_with_base:
            from peft import PeftModel

            base_model_id = _get_base_model_id(job_id)
            if base_model_id:
                device_map = {"": 0}

                base_model = AutoModelForCausalLM.from_pretrained(
                    base_model_id, torch_dtype=torch.float16, device_map=device_map
                )
                peft_model = PeftModel.from_pretrained(base_model, adapter_path)
                merged = peft_model.merge_and_unload()
                merged.save_pretrained(output_path)
                tokenizer = AutoTokenizer.from_pretrained(adapter_path)
                tokenizer.save_pretrained(output_path)
            else:
                from peft import PeftModel

                merged = PeftModel.from_pretrained(
                    AutoModelForCausalLM.from_pretrained(
                        str(adapter_path), torch_dtype=torch.float16, device_map={"": 0}
                    ),
                    adapter_path,
                )
                merged = merged.merge_and_unload()
                merged.save_pretrained(output_path)
                tokenizer = AutoTokenizer.from_pretrained(adapter_path)
                tokenizer.save_pretrained(output_path)
        else:
            shutil.copytree(adapter_path, output_path)

        result = {"output_path": str(output_path)}

        if export_format == "gguf":
            gguf_path = output_path.with_suffix(".gguf")
            _export_to_gguf(output_path, gguf_path)
            result["gguf_path"] = str(gguf_path)

        return result

    except Exception as e:
        import traceback
        return {"error": str(e), "traceback": traceback.format_exc()}


def _get_base_model_id(job_id: str):
    config_path = _ADAPTERS_DIR / job_id / "adapter_config.json"
    if config_path.exists():
        with open(config_path) as f:
            return json.load(f).get("base_model_name_or_path")
    return None


def _export_to_gguf(model_path: Path, output_path: Path):
    try:
        import subprocess
        result = subprocess.run(
            ["python", "-m", "llama_cpp.convert_hf_to_gguf", str(model_path), "--outfile", str(output_path)],
            capture_output=True, text=True, timeout=600,
        )
        if result.returncode != 0:
            raise RuntimeError(result.stderr)
    except Exception as e:
        raise RuntimeError(f"GGUF export not available: {e}")