from dataclasses import dataclass
from pathlib import Path
import subprocess

from varidock.types import Trajectory, ConformationSet, PDB
from varidock.pipeline.stage import Stage
from varidock.utils import run_with_interrupt

@dataclass
class VMDFrameExtractionConfig:
    output_dir: Path

class VMDFrameExtraction(Stage[Trajectory, ConformationSet]):
    name = "vmd_frame_extraction"
    input_type = Trajectory
    output_type = ConformationSet

    def __init__(self, config: VMDFrameExtractionConfig):
        self.config = config

    def run(self, input: Trajectory) -> ConformationSet:
        # 1. Load PSF
        # 2. Load each .coor file as a frame
        # 3. Select protein atoms
        # 4. Write each frame as PDB
        # 5. Return ConformationSet with list of PDBs
        self.config.output_dir = self.config.output_dir.resolve()
        self.config.output_dir.mkdir(parents=True, exist_ok=True)

        tcl = f"""
            set molid [mol new "{input.psf.path.resolve()}"]
            mol addfile {input.pdb.path.resolve()} $molid
            """
        
        for coor in input.coor_files:
            tcl += f'mol addfile "{coor.resolve()}" $molid\n'
        
        tcl += f"""
        package require topotools
        set prot [atomselect $molid "protein"]
        topo -molid $molid -sel $prot guessatom element mass            

        $prot writepsf {self.config.output_dir}/protein.psf

        set numframes [molinfo top get numframes]
        for {{set i 0}} {{$i < $numframes}} {{incr i}} {{
            $prot frame $i
            $prot writepdb {self.config.output_dir}/protein_conf$i.pdb
        }}
        $prot delete
        mol delete $molid
        exit
        """

        script_path = self.config.output_dir / "frame_extraction.tcl"
        script_path.parent.mkdir(parents=True, exist_ok=True)
        script_path.write_text(tcl)

        with open(self.config.output_dir / "vmd_frame_extraction.log", "w") as f:
            run_with_interrupt(
                ["vmd", "-dispdev", "none", "-eofexit", "-e", str(script_path)],
                stdout=f,
                stderr=subprocess.STDOUT,
                cwd=self.config.output_dir,
            )
            
        pdb_files = sorted(self.config.output_dir.glob("protein_conf*.pdb"))
        return ConformationSet(
            psf=self.config.output_dir / "protein.psf",
            pdbs=[PDB(path=p) for p in pdb_files],
        )
    




