import Foundation
import Path

public extension String {
    func findFirst(pat: String) -> Range<String.Index>? {
        return range(of: pat, options: .regularExpression)
    }
    func hasMatch(pat: String) -> Bool {
        return findFirst(pat:pat) != nil
    }
    func withMaj() -> String {
        return prefix(1).capitalized + dropFirst()
    }
}

public func nbNameToScriptName(_ name: String) -> String {
    var splits = name.components(separatedBy: "_")
    splits = splits[1...].map { $0.withMaj() }
    return splits.joined(separator: "")
}

public func readNb(_ fname: Path) -> [String:Any] {
    let data = try! Data(contentsOf: fname.url)
    return try! JSONSerialization.jsonObject(with: data, options: .allowFragments) as! [String: Any]
}

public func notebookToScript(_ fname: Path, dest: Path?=nil) {
    let newName = nbNameToScriptName(fname.basename(dropExtension: true))+".swift"
    let destFname = (dest ?? fname.parent) / newName
    let cells = readNb(fname)["cells"] as! [[String:Any]]
        var module = """
/*
This file was autogenerated from \(fname.basename())
        
If you edit it, be sure that:
  1. there is no diff between this file and the corresponding notebook prior to editing
  2. you don't touch the comments looking like // cell ## as it would break the way back to the notebook
        
Run *** when you are done to update the notebooks with your change.
*/
        
"""
    for (i,cell) in cells.enumerated() {
        if let source = cell["source"] as? [String], !source.isEmpty, source[0].hasMatch(pat: #"^\s*//\s*export\s*$"#) {
            module.append("\n//cell\(i)\n\(source[1...].joined())\n")
        }
    }
    try! module.write(to: destFname, encoding: .utf8)
}

let fname = path/"/nbs/stable/00_load_data.ipynb"
let dest = path/"Sources/SwiftAI"

notebookToScript(fname, dest: dest)

public func makeLibrary(_ nbFolder: Path, dest: Path?=nil){
    for entry in try! nbFolder.ls()  where entry.kind == Entry.Kind.file && entry.path.basename().hasMatch(pat: #"^\d+[a-z]*_.*ipynb$"#) {
        print("Converting \(entry.path.basename())")
        notebookToScript(entry.path, dest: dest)
    }
}

let nbFolder = path/"/nbs/stable"

let fname = dest/"LoadData.swift"

public func readScript(_ fname: Path) -> String {
    let data = try! Data(contentsOf: fname.url)
    return String(data: data, encoding: .utf8)!
}

public func writeNotebook(_ nbFname: Path, nbData: [String: Any]) {
    let outData = try! JSONSerialization.data(withJSONObject: nbData, options: .prettyPrinted)  
    let jsonString = String(data: outData, encoding: .utf8)! 
    do { try jsonString.write(to: nbFname, encoding: .utf8) }
    catch { "Couldn't save notebook" }
}

public func scriptToNotebook(_ fname: Path, nbFolder: Path){
    let code = readScript(fname)
    let codeCells = code.components(separatedBy: "//cell")
    let nbName = codeCells[0].components(separatedBy: "\n")[1].components(separatedBy: "from ")[1]
    let nbFname = nbFolder / nbName
    var jsonData = readNb(nbFname)
    var cells = jsonData["cells"] as! [[String:Any]]
    for c in codeCells[1...] {
        var lines = c.components(separatedBy: "\n")
        let idx: Int = Int(lines[0])!
        var i = lines.count-1
        while lines[i].isEmpty { i -= 1}
        if i > 1 {
            for i in 1...(i-1) { lines[i].append("\n") }
        }
        lines[0] = "// export\n"
        cells[idx]["source"] = Array(lines[...i])
    }
    jsonData["cells"] = cells
    writeNotebook(nbFname, nbData: jsonData)
}

scriptToNotebook(fname, nbFolder: nbFolder)

public func updateNotebooks(_ scriptsFolder: Path, nbFolder: Path) {
    for entry in try! scriptsFolder.ls()  where entry.kind == Entry.Kind.file && entry.path.basename().hasMatch(pat: #".swift$"#) {
        print("Updating nb from \(entry.path.basename())")
        scriptToNotebook(entry.path, nbFolder: nbFolder)
    }
}

public func diffNbScript(_ nbFname: Path, dest: Path){
    let newName = nbNameToScriptName(fname.basename(dropExtension: true))+".swift"
    let destFname = (dest ?? fname.parent) / newName
    let cells = readNb(nbFname)["cells"] as! [[String:Any]]
    
    let data = try! Data(contentsOf: fname.url)
    let code: String = String(data: data, encoding: .utf8)!
    let codeCells = code.components(separatedBy: "//cell")
    let nbName = codeCells[0].components(separatedBy: "\n")[1].components(separatedBy: "from ")[1]
    let nbData = try! Data(contentsOf: nbFname.url)
    var jsonData = try! JSONSerialization.jsonObject(with: nbData, options: .allowFragments) as! [String: Any]
    var cells = jsonData["cells"] as! [[String:Any]]
    for c in codeCells[1...] {
        var lines = c.components(separatedBy: "\n")
        let idx: Int = Int(lines[0])!
        var i = lines.count-1
        while lines[i].isEmpty { i -= 1}
        if i > 1 {
            for i in 1...(i-1) { lines[i].append("\n") }
        }
        lines[0] = "// export\n"
        cells[idx]["source"] = Array(lines[...i])
    }
    jsonData["cells"] = cells
    let outData = try! JSONSerialization.data(withJSONObject: jsonData, options: .prettyPrinted)  
    let jsonString = String(data: outData, encoding: .utf8)! 
    do { try jsonString.write(to: nbFname, encoding: .utf8) }
    catch { "Couldn't save notebook" }
}