import shutil
from pathlib import Path
from ase.io import read, write
from nep_auto.modules.base_module import BaseModule


class SCFModule(BaseModule):
    def __init__(self, driver, iter_id):
        super().__init__(driver, iter_id)
        self.template_subdir = "02_scf"
        self.work_dir = self.iter_dir / "02.scf"
        self.select_dir = self.iter_dir / "01.select"

    def get_work_dir(self):
        return self.work_dir

    def run(self):
        self.logger.info(f"⚛️ [SCF] Starting DFT Calculation Iter {self.iter_id}...")
        self.initialize()

        # 1. 读取 selected.xyz
        selected_xyz = self.select_dir / "selected.xyz"
        if not selected_xyz.exists():
            raise FileNotFoundError("selected.xyz missing")

        self.logger.info("   -> Reading structures using ASE...")
        atoms_list = read(selected_xyz, index=':')
        self.logger.info(f"   -> Found {len(atoms_list)} structures.")

        # 2. 准备任务文件夹
        task_dirs = []
        for i, atoms in enumerate(atoms_list):
            task_name = f"task.{i:03d}"
            task_dir = self.work_dir / task_name
            task_dir.mkdir(exist_ok=True)
            task_dirs.append(task_dir)

            # 写 POSCAR
            write(task_dir / "POSCAR", atoms, format='vasp')

            # 复制模版 INCAR, KPOINTS, POTCAR
            self.copy_template("INCAR", target_name=None)  # 复制到 self.work_dir
            shutil.copy(self.work_dir / "INCAR", task_dir / "INCAR")  # 再分发
            self.copy_template("KPOINTS", target_name=None)
            shutil.copy(self.work_dir / "KPOINTS", task_dir / "KPOINTS")
            self.copy_template("POTCAR", target_name=None)
            shutil.copy(self.work_dir / "POTCAR", task_dir / "POTCAR")

        # 3. 提交任务
        # 这里区分 local 模式和 slurm 模式
        # 既然你目前是 interactive gpu，我们假设是串行或者简单的并行
        self.logger.info("   -> Running VASP jobs...")

        success_count = 0
        for task_dir in task_dirs:
            self.logger.info(f"      -> Running {task_dir.name}...")
            try:
                # 调用 machine.yaml 里定义的 vasp
                # 注意：如果 task 很多，这里最好写成多进程并发
                self.runner.run("vasp", cwd=task_dir)

                # 简单检查
                if (task_dir / "OUTCAR").exists():
                    success_count += 1
            except Exception as e:
                self.logger.error(f"Task {task_dir.name} failed: {e}")

        self.logger.info(f"   -> Finished. Success: {success_count}/{len(task_dirs)}")

        # 4. 收集数据 (OUTCAR -> NEP-dataset.xyz)
        self.logger.info("   -> Collecting data...")
        valid_atoms = []
        for task_dir in task_dirs:
            try:
                # 读取 OUTCAR
                atoms = read(task_dir / "OUTCAR", format='vasp-outcar')
                valid_atoms.append(atoms)
            except:
                pass

        if valid_atoms:
            write(self.work_dir / "NEP-dataset.xyz", valid_atoms, format='extxyz')
        else:
            raise RuntimeError("No valid OUTCARs found!")

        self.check_done()

    def check_done(self):
        if (self.work_dir / "NEP-dataset.xyz").exists():
            return True
        raise RuntimeError("SCF failed: NEP-dataset.xyz not generated")