import pya

import sys
import re # regular expressions
import os
import pathlib
import time
from pathlib import Path
import threading
import queue


# here defines the work location
workpath = "C:/Users/jonathan.mas/Documents/klayout_frametest/"
# in the work folder two subfolders are required:
Lfolder = workpath + 'Libraries/' # libraries
Efolder = workpath + 'Export/' # export
# max level of hierarchy to view:
maxOpLvl = 1000
# example: If you set  maxOpLvl= 3, when you open a design you will never get more than 3 level of hierarchy displayed
# The level of hierachy displayed is betermined by:
# 1- 5second worker wait
# 2- nevermore than maxOpLvl
# Note: If you open the same layout at a later time you will get more hierarchy, as the layout is fully opening in a thread the moment you are opening it
# Note2: If you edit and save a cell in the hierarchy the loader will be stopped
gLiblist = [] # global list of the availables libraries
gCellList = {} # global ditionary of the cells in libraries
GlbSnap = 1 # snap (um) to use for tools like insert cell
OpnLibs = [] # list of opened libs
# libraries are opened:
# when you click to get the listing in the library manager
# when you open/insert a design with a reference to it
# they stay opened during the whole session
OpnLays = []
# list open layouts [cellname,libname,lvl, loader]
# loader is the thread in which it is getting loaded
# lvl is for the level in the loader if lvl is 0 it is a topcell in the loader

insertLibCell = [] # global variable to the latest cell that has been inserted with the insert tool


def load_Libs(directory):
  liblist = []
  for file_name in os.listdir(directory):
    file_path = os.path.join(directory, file_name)
    if os.path.isfile(file_path) and file_name.endswith('.oas'):  # Only add files, not directories
      liblist.append(os.path.splitext(file_name)[0])
  return liblist

gLiblist = load_Libs(Lfolder)

def load_Cells(libname, usl=True): #will make a cell list for the selected library
  if usl and (libname in gCellList) : return gCellList[libname]
  ly = pya.Layout()
  #start_time = time.time()
  libindex = getLibIndex(libname)
  #end_time = time.time()
  #elapsed_time = end_time - start_time
  #print(f"Time taken: {elapsed_time:.2f} seconds")
  gCellList[libname]=[]
  for cell in OpnLibs[libindex][0].top_cells():
    gCellList[libname].append(cell.name)
  return gCellList[libname]



class Libman(pya.QMainWindow):
  def __init__(self):
    super().__init__()
    self.setWindowTitle("library mannager")
    self.setGeometry(100, 100, 600, 400)    # Create the menu bar
    menu_bar = self.menuBar
    file_menu = menu_bar.addMenu("File")    # Add actions to the menu
    # new library menu
    self.NewLib = pya.QAction("New Library", self)
    self.NewLib.triggered(lambda event : self.newlib())
    file_menu.addAction(self.NewLib)    
    # delete library menu
    self.DelLib = pya.QAction("delete Library", self)
    self.DelLib.triggered(lambda event : self.dellib())
    file_menu.addAction(self.DelLib)
    # import menu
    self.ImportFile = pya.QAction("Import", self)
    self.ImportFile.triggered(lambda event : self.importFile())
    file_menu.addAction(self.ImportFile)
    # export menu
    self.ExportFile = pya.QAction("Export", self)
    self.ExportFile.triggered(lambda event : self.exportFile())
    file_menu.addAction(self.ExportFile)
    # create a new cell
    self.NC = pya.QAction("New Cell", self)
    self.NC.triggered(lambda event : self.newcell())
    file_menu.addAction(self.NC)
     # create a new cell
    self.DC = pya.QAction("Delete Cell", self)
    self.DC.triggered(lambda event : self.delcell())
    file_menu.addAction(self.DC)
    
    self.mw = pya.Application.instance().main_window()
    self.mw.create_view()
    
    self.wi = pya.QWidget()
    self.setCentralWidget(self.wi)   
    self.libraries = pya.QListWidget()
    self.cells = pya.QListWidget()    # Add sample items to the list boxes
    self.cells.setSortingEnabled(True)
    self.refresh_libs()
    self.libraries.itemClicked(lambda item : self.refresh_cells(item))
    self.cells.itemDoubleClicked(lambda item : self.open_mvl_Cell(item))
    layout = pya.QHBoxLayout()
    layout.addWidget(self.libraries)
    layout.addWidget(self.cells)    # Set the layout for the central widget
    self.wi.setLayout(layout)

  def refresh_libs(self):
    self.libraries.clear()  # Clear the listbox
    self.libraries.addItems(gLiblist)
    
  def refresh_cells(self, item): #will make a cell list for the selected library
    self.cells.clear()
    cell_list = load_Cells(item.text)
    self.cells.addItems(cell_list)

  def open_mvl_Cell(self, item):
    global insertLibCell
    liblistIpt = self.libraries.selectedItems()
    if liblistIpt:
      start_time = time.time()
      mainlayout = pya.Layout()
      place_Cell(mainlayout, item.text, liblistIpt[0].text)
      self.mw.current_view().show_layout(mainlayout, False)
      # the previous layout will be destroyed not bad but if we want to implement a layout save later/ comeback later
      # we would need to:
      # 1 -create a list of Worked Layouts:
      # 2- duplicate the layout when another layout is called => if ly== self.mw.current_view().active_cellview().layout(): ly=ly.dup()
      # there is also a new layout handle develped from version 0.30.6 that can be usefull
      elapsed_time = time.time() - start_time
      print(f"Time taken to open {item.text}: {elapsed_time:.2f} seconds")
      if not insertLibCell:
        insertLibCell = [liblistIpt[0].text,item.text]
    

  def newlib(self):
    libname = pya.QInputDialog.getText(None, "Input Library name", "Library name:")
    if libname:
      ly = pya.Layout()
      ly.write(Lfolder + libname + ".oas")
      ly.clear()
      gLiblist.append(libname)
      self.refresh_libs()

  def newcell(self):
    libname = self.libraries.itemFromIndex(self.libraries.selectedIndexes()[0]).text
    cellname = pya.QInputDialog.getText(None, "Input New Cell name", "New cell in " + libname + ":")
    libindex = getLibIndex(libname)
    if cellname:
      OpnLibs[libindex][0].create_cell(cellname)
      OpnLibs[libindex][0].write(Lfolder + libname + ".oas")
      self.cells.addItem(cellname)
      gCellList[libname].append(cellname)
      

  def delcell(self):
    libname = self.libraries.itemFromIndex(self.libraries.selectedIndexes()[0]).text
    cellname = self.cells.itemFromIndex(self.cells.selectedIndexes()[0]).text
    libindex = getLibIndex(libname)
    if cellname:
      OpnLibs[libindex][0].delete_cell(OpnLibs[libindex][0].cell(cellname).cell_index())
      OpnLibs[libindex][0].write(Lfolder + libname + ".oas")
      self.cells.takeItem(self.cells.selectedIndexes()[0].row())
      # needs to delete the hierarchy too?


  def dellib(self):
    liblistpt = self.libraries.selectedIndexes()
    if liblistpt:
      os.remove(Lfolder + self.libraries.itemFromIndex(liblistpt[0]).text + ".oas")
      gLiblist.remove(self.libraries.itemFromIndex(liblistpt[0]).text)
      self.refresh_libs()

  def importFile(self):
    filename = pya.QFileDialog.getOpenFileName(None, "File To import", workpath, "desgnfiles (*.gds *.oas *.gds.gz *.oas.gz);;All Files (*)")
    if not filename :
      return
    
    libname = pathlib.Path(filename).stem.split('.')[0]
    start_time = time.time()
    logfile = open(workpath + "import.log", "w")
    ly = pya.Layout()
    ly.read(filename)
    boxlayerIndex = ly.layer(1002, 0)
    cell_dict = {}
    print(f"Now import {libname}\n")
    for importcell in ly.each_cell():
      logfile.write(f"import {importcell.name}\n")
      for inst in importcell.each_inst():
        instname= inst.cell.name + "!>" + libname
        if instname in cell_dict:
          ptc = ly.cell(instname)
          logfile.write(f"  Link instance {instname}\n")
        else :
          ptc = ly.create_cell(instname)
          logfile.write(f"  create instance {instname}\n")
          box = pya.Box(inst.cell.bbox())
          ptc.shapes(boxlayerIndex).insert(box)
        cell_dict[instname] = True
        importcell.insert(pya.CellInstArray(ptc.cell_index(), inst.cplx_trans, inst.a, inst.b, inst.na, inst.nb))
        inst.delete()
      
    ly.write(Lfolder + libname + ".oas")
    gLiblist.append(libname)
    self.refresh_libs()
    elapsed_time = time.time() - start_time
    print(f"Time taken_import {libname}: {elapsed_time:.2f} seconds\n")


  def exportFile(self):
    liblistIpt = self.libraries.selectedItems()
    celllistIpt = self.cells.selectedItems()
    if liblistIpt and celllistIpt:
      start_time = time.time()
      cellname= celllistIpt[0].text
      libname = liblistIpt[0].text
      try:
        index = next(i for i, ly in enumerate(OpnLays) if (ly[0] == cellname) and (ly[1] == libname) and ly[2]==0)
      except StopIteration:
        index = len(OpnLays)
        loader = CellLoader(1)
        loader.load(cellname, libname)
        OpnLays.append([cellname,libname,0, loader])
      ly = OpnLays[index][3].get_mainlayout()
      for exportcell in ly.each_cell():
        if "!>" not in exportcell.name:
          continue
        cn, ln = exportcell.name.split("!>")
        new_name = cn
        index =0
        while ly.cell(new_name) is not None:
          index += 1
          new_name = f"{cn}_{index}"
        exportcell.name=new_name
      lo = pya.SaveLayoutOptions()
      lo.format= "OASIS"
      ly.write(Efolder + cellname +".oas",lo)
      elapsed_time = time.time() - start_time
      print(f"Time taken_export : {elapsed_time:.2f} seconds\n")

libman = Libman()
libman.show()

class saveCell (pya.Action):
  def triggered(self):
      # needs to clear cellview loader when saving
      view = pya.Application.instance().main_window().current_view()
      myly = view.active_cellview().layout()
      sc = myly.top_cell()
      cln = sc.name
      cn, ln = cln.split("!>")
      libindex = getLibIndex(ln)
      boxlayerIndex = OpnLibs[libindex][0].layer(1002, 0)
      tc = OpnLibs[libindex][0].cell(cn)
      tcCnL = [inst.cell.name for inst in tc.each_inst()]
      tc.clear()
      tc.copy_shapes(sc)
      # instances save 
      for inst in sc.each_inst():
        tcln = inst.cell.name
        mc=OpnLibs[libindex][0].cell(tcln)
        if not mc:
          mc = OpnLibs[libindex][0].create_cell(tcln)
          box = pya.Box(inst.cell.bbox())
          mc.shapes(boxlayerIndex).insert(box)
        tc.insert(pya.CellInstArray(mc.cell_index(),inst.cplx_trans, inst.a, inst.b, inst.na, inst.nb))
      for tcn in tcCnL:
        ptcn = OpnLibs[libindex][0].cell(tcn)
        if ptcn and ptcn.is_top():
          OpnLibs[libindex][0].delete_cell(ptcn.cell_index())
      OpnLibs[libindex][0].write(Lfolder + OpnLibs[libindex][1] + ".oas")
      clearLoaders(cn, ln)

save_handler = saveCell()
pya.Application.instance().main_window().menu().insert_item('@toolbar.end', 'Save', save_handler)
save_handler.title = "Save"


class insertInstance(pya.Plugin):
  def __init__(self):
    super(insertInstance, self).__init__()
    pass
  def activated(self):
    self.grab_mouse()
    self.layH =  pya.Application.instance().main_window().current_view().active_cellview().layout().layer(10000, 0)
    pya.Application.instance().main_window().current_view().add_missing_layers()
    layer = MySubs.GetLayer(10000,0)
    layer.visible = True
    layer.name = "Marker_Layer"
    self.top = pya.Application.instance().main_window().current_view().active_cellview().layout().top_cell()
    self.dialog = pya.QDialog(pya.Application.instance().main_window())
    QLay = pya.QGridLayout(self.dialog)
    self.libs =  pya.QComboBox()
    self.cells =  pya.QComboBox()
    self.libs.addItems(gLiblist)
    self.libs.currentText= insertLibCell[0]
    self.update_cells()
    self.cells.currentText= insertLibCell[1]
    self.bbx = [0,0,2000,2000]
    self.cn =  self.cells.currentText + "!>" + self.libs.currentText
    self.setCellparam()
    self.libs.currentTextChanged(self.update_cells)
    self.cells.currentTextChanged(self.setCellparam)
    
    l1 =  pya.QLabel('snap (um)= ',self.dialog)
    self.l2 = pya.QLineEdit('{0:g}'.format(GlbSnap),self.dialog)
    self.l2.setToolTip('this will snap to grid')
    b1= pya.QPushButton('apply!',self.dialog)
    b1.clicked = self.ApplySnap
    QLay.addWidget(self.libs,1,0)
    QLay.addWidget(self.cells,2,0)
    QLay.addWidget(l1,3,0)
    QLay.addWidget(self.l2,3,1)
    QLay.addWidget(b1,4,1)
    self.dialog.show()
  def update_cells(self):
    self.cells.clear()  # Remove old items
    if not self.libs.currentText in gCellList:
      load_Cells(self.libs.currentText)
    self.cells.addItems(gCellList[self.libs.currentText]) 
    insertLibCell[0] = self.libs.currentText   
    self.setCellparam()
  def setCellparam(self):
    if self.cells.currentText and self.libs.currentText:
      self.cn =  self.cells.currentText + "!>" + self.libs.currentText
      libindex = getLibIndex(self.libs.currentText)
      self.bbx = OpnLibs[libindex][0].cell(self.cells.currentText).bbox()
      insertLibCell[1] = self.cells.currentText
      print(f"Set insertLibCell to {insertLibCell}\n") 
  def ApplySnap(self) :
    global GlbSnap
    GlbSnap= float(self.l2.text)
  def deactivated(self):
    pya.Application.instance().main_window().current_view().active_cellview().layout().clear_layer(self.layH)
    pya.Application.instance().main_window().current_view().remove_unused_layers()
    pya.MainWindow.instance().message("", 0)
    self.ungrab_mouse()
    self.dialog.close()
  def mouse_moved_event(self, p, buttons, prio):
    if prio:
      self.set_cursor(pya.Cursor.Blank)
      pya.Application.instance().main_window().current_view().active_cellview().layout().clear_layer(self.layH)
      xMouse= int(p.x/GlbSnap)*GlbSnap
      yMouse= int(p.y/GlbSnap)*GlbSnap
      text = pya.Text('x='+str(xMouse)+' y='+str(yMouse), xMouse*1000,yMouse*1000)
      text.halign = 1
      text.valign = 2
      self.top.shapes(self.layH).insert(text)
      rect =pya.Box(xMouse*1000+self.bbx.p1.x-1,yMouse*1000+self.bbx.p1.y-1,xMouse*1000+self.bbx.p2.x,yMouse*1000+self.bbx.p2.y)
      self.top.shapes(self.layH).insert(rect)
      return False
  def mouse_click_event(self, p, buttons, prio):
    if prio:
      view = pya.Application.instance().main_window().current_view()
      layout = view.active_cellview().layout()
      tc = layout.top_cell()
      xMouse= int(p.x/GlbSnap)*GlbSnap
      yMouse= int(p.y/GlbSnap)*GlbSnap
      if tc.name == self.cn:
        pya.QMessageBox.warning(None, 'Instance error', 'circular references are not allowed')
      else:
        place_Cell(layout, self.cells.currentText,self.libs.currentText)
        tc.insert(pya.CellInstArray(layout.cell(self.cn).cell_index(), pya.Trans(0, False, xMouse*1000, yMouse*1000)))
      self.ungrab_mouse()
      pya.MainWindow.instance().cancel()
      return True
    return False

class insertInstanceFactory(pya.PluginFactory):
  def __init__(self):
    super(insertInstanceFactory, self).__init__()
    pya.MainWindow.instance()
    TheButton = self.register(-1000, "Insert_cell", "Insert")
  def create_plugin(self, manager, root, view):
    return insertInstance()
  instance = None
insertInstanceFactory.instance = insertInstanceFactory()

def getLibIndex(ln):
  try:
    libindex = next(i for i, lib in enumerate(OpnLibs) if lib[1] == ln)
  except StopIteration:
    libindex = len(OpnLibs)
    OpnLibs.append([pya.Layout(), ln])
    OpnLibs[libindex][0].read(Lfolder + ln + ".oas")
  return libindex

def place_Cell(mainlayout, cellname, libname):
  timer =0
  try:
    index = next(i for i, ly in enumerate(OpnLays) if (ly[0] == cellname) and (ly[1] == libname) and (ly[2]==0 or ly[3].current_level-ly[2]>4 or not ly[3].is_running()))
  except StopIteration:
    index = len(OpnLays)
    #locCelllayout = pya.Layout()
    loader = CellLoader(maxOpLvl)
    loader.load(cellname, libname)
    OpnLays.append([cellname,libname,0, loader])
    timer = 5
  cn =  cellname + "!>" + libname
  loadlayout = OpnLays[index][3].get_snapshot(timer)
  sc = loadlayout.cell(cn)
  tc = mainlayout.create_cell(cn)
  tc.copy_tree(sc)
  loadlayout.clear()

def clearLoaders(cellname,libname):
  for oly in OpnLays[:]:
    if oly[0]==cellname and oly[1]==libname:
      oly[3].stop()
      OpnLays.remove(oly)
  
class CellLoader:
  def __init__(self, maxlvl=5):
    self.mainlayout = pya.Layout()
    #self.snapshot_timeout = snapshot_timeout
    self.exported_cells = []
    self._lock = threading.Lock()
    self._result_queue = queue.Queue()
    self._worker = None
    self._ic_list = None
    self.current_level = 0
    self.snapLayout = self.mainlayout.dup()
    self.finishedsnap = False
    self.maxlvl = maxlvl
    self._stop_event = threading.Event()
  # ------------------------------------------------------------------ #
  # load a design in a thread
  # when you start 
  # loader = CellLoader(maxlvl)
  # maxlvl wil define the max snapshot level
  # loader.load(cellname, libname)
  # it will start a thread loading a design inside  loader.mainlayout
  # you can snapshots with loader.get_snapshot(snapshot_timeout)
  # if the worker is still running, it will wait snapshot_timeout seconds to see if it complete, before starting the snapshot process
  # snapshot will stop at maxlvl could be a usefull option if we want very fast design loading
  # ------------------------------------------------------------------ #

  def load(self, top_cell_name, lib_name):
    self._stop_event.clear()
    self._ic_list = self._init_top_cell(top_cell_name, lib_name)
    self._worker = threading.Thread(
        target=self._loadcell,
        args=(self._ic_list, 0),
        daemon=True
    )
    self._worker.start()

  def stop(self):
    self._stop_event.set()
    if self._worker is not None:
      self._worker.join()  # wait for the thread to fully exit
    with self._lock:
      self.mainlayout.clear()
      self.snapLayout.clear()
      self.exported_cells = []   # optional but avoids stale cache on reload
      self._ic_list = None
        
  def get_snapshot(self,snapshot_timeout):
    if self._worker is None:
      raise RuntimeError("Call load() first.")
    self._worker.join(timeout=snapshot_timeout)
    if self._worker.is_alive() or self.current_level > self.maxlvl:
      with self._lock:
        if not self.finishedsnap:
          self._finishbbx()
          self.finishedsnap = True
        #print("Snapshot done with box completion")
        return self.snapLayout.dup()
    else:
      self.snapLayout.clear()
      return self.mainlayout.dup()
      print("Worker already finished")
      #snapLayout = self.mainlayout

  def get_mainlayout(self):
    if self._worker is None:
      raise RuntimeError("Call load() first.")
    self._worker.join()  # blocks until worker finishes
    return self.mainlayout
  # ------------------------------------------------------------------ #
  #  
  # ------------------------------------------------------------------ #

  def _init_top_cell(self, top_cell_name, lib_name):
      lib_index = getLibIndex(lib_name)
      target_cell = self.mainlayout.create_cell(top_cell_name + "!>" + lib_name)
      source_cell = OpnLibs[lib_index][0].cell(top_cell_name)
      return self._copy_cell_and_collect_instances(source_cell, target_cell)

  def _copy_cell_and_collect_instances(self, source_cell, target_cell):
      target_cell.copy_shapes(source_cell)
      return [[inst, target_cell] for inst in source_cell.each_inst()]

  def _loadcell(self, ic_list, start_lvl):
      while ic_list:
        if self._stop_event.is_set():
          break
        self.current_level += 1
        new_ic_list = []
        for ic in ic_list:
          if self._stop_event.is_set():
            break
          cln = ic[0].cell.name
          if "!>" in cln:
            cn, ln = cln.split("!>", 1)
            if cln not in self.exported_cells:
              lib_index = getLibIndex(ln)
              tc = self.mainlayout.create_cell(cn + "!>" + ln)
              sc = OpnLibs[lib_index][0].cell(cn)
              new_ic_list += self._copy_cell_and_collect_instances(sc, tc)
              self.exported_cells.append(cln)
              OpnLays.append([cn,ln,self.current_level, self])
            ic[1].insert(pya.CellInstArray(
              self.mainlayout.cell(cln).cell_index(),
              ic[0].cplx_trans, ic[0].a, ic[0].b, ic[0].na, ic[0].nb
            ))
        ic_list = new_ic_list
        with self._lock:
          if self.current_level < self.maxlvl+1:
            self.snapLayout.clear()
            self.snapLayout = self.mainlayout.dup()
            self.finishedsnap = False
            self._ic_list= new_ic_list
      if self._stop_event.is_set():
        self._result_queue.put(None)
      else:
        self._result_queue.put(self.mainlayout)

  def _finishbbx(self):
      #Insert bounding boxes for all instances in self._ic_list
      box_layer_index = self.snapLayout.layer(1002, 0)
      for ic in self._ic_list:
          box = pya.Box(ic[0].bbox())
          self.snapLayout.cell(ic[1].name).shapes(box_layer_index).insert(box)
  
  def is_running(self):
    return self._worker is not None and self._worker.is_alive()
    #if self._worker is None or self._worker.is_alive(): return True 
    #else: return False