import Foundation


public struct ObjReader
{
	public static func read(url: URL) throws -> ObjModel
	{
		var file = try ObjDocumentReader(filePath: url)

		var model = ObjModel()
		var name: String? = nil
		var object = ObjModel.Object()
		var mesh = ObjModel.Mesh()

		file.string(tag: "mtllib") { s in
			let mats = try ObjMtlLoader.read(url: url.deletingLastPathComponent().appendingPathComponent(s[0]))
			model.materials.merge(mats, uniquingKeysWith: { (_, new) in new } )
		}
		file.string(tag: "o", count: 1) { s in
			if !mesh.faces.isEmpty { object.meshes.append(mesh) }
			if !object.meshes.isEmpty
			{
				model.objects[name!] = object
				object = .init()
				mesh = .init()
			}
			name = String(s[0])
		}
		file.float(tag: "v") { v in
			if v.count == 6
			{
				if model.colours.isEmpty && !model.positions.isEmpty
				{
					for _ in 0..<model.positions.count
					{
						model.colours.append(.one)
					}
				}
				model.colours.append(Vec3f(v[3], v[4], v[5]))
			}
			else if !model.colours.isEmpty && v.count == 3
			{
				model.colours.append(.one)
			}
			else if v.count != 3 { throw ObjLoaderError.badTagParameters }
			model.positions.append(Vec3f(v[0], v[1], v[2]))
		}
		file.float(tag: "vn", count: 3) { v in model.normals.append(Vec3f(v[0], v[1], v[2])) }
		file.float(tag: "vt", count: 2) { v in model.texCoords.append(Vec2f(v[0], v[1])) }
		file.int(tag: "l") { i in mesh.faces.append(.line(p1: i[0], p2: i[1])) }
		file.raw(tag: "f") { raw in
			let readIndex =
			{ (raw: Substring) in
				let face = raw.split(separator: "/")
				guard face.count == 3,
					let posIdx = Int(face[0]), let coordIdx = Int(face[1]), let normIdx = Int(face[2])
				else { throw ObjLoaderError.badTagParameters }
				return ObjModel.Index(p: posIdx - 1, n: normIdx - 1, t: coordIdx - 1)
			}

			if raw.count == 3
			{
				for raw in raw { _ = try readIndex(raw) }
				mesh.faces.append(.triangle(
					v1: try readIndex(raw[raw.startIndex]),
					v2: try readIndex(raw[raw.startIndex + 1]),
					v3: try readIndex(raw[raw.startIndex + 2])))
			}
			else if raw.count == 4
			{
				mesh.faces.append(.quad(
					v1: try readIndex(raw[raw.startIndex]),
					v2: try readIndex(raw[raw.startIndex + 1]),
					v3: try readIndex(raw[raw.startIndex + 2]),
					v4: try readIndex(raw[raw.startIndex + 3])))
			}
			else if raw.count >= 5
			{
				mesh.faces.append(.ngon(try raw.map { try readIndex($0) }))
			}
			else { throw ObjLoaderError.badTagParameters }
		}
		file.string(tag: "usemtl", count: 1) { s in
			if !mesh.faces.isEmpty
			{
				object.meshes.append(mesh)
				mesh = .init()
			}
			mesh.material = s[0]
		}
		// file.int(tag: "s", count: 1) { i in } //TODO: Smoothing groups

		try file.read()

		if !mesh.faces.isEmpty { object.meshes.append(mesh) }
		if !object.meshes.isEmpty { model.objects[name!] = object }

		return model
	}
}

fileprivate struct ObjMtlLoader
{
	static func read(url: URL) throws -> Dictionary<String, ObjMaterial>
	{
		var file = try ObjDocumentReader(filePath: url)

		var materials = Dictionary<String, ObjMaterial>()
		var name: String = ""
		var mat: ObjMaterial? = nil

		file.string(tag: "newmtl", count: 1)
		{ s in
			if mat != nil
			{
				materials[name] = mat!
			}
			mat = .init()
			name = String(s[0])
		}

		file.preHandle { s in if s != "newmtl" && mat == nil { throw ObjLoaderError.unexpectedTag } }
		file.float(tag: "Ka", count: 3) { f in mat!.ambient = Colour(r: f[0], g: f[1], b: f[2]) }  // "Ambient colour"
		file.float(tag: "Kd", count: 3) { f in mat!.diffuse = Colour(r: f[0], g: f[1], b: f[2]) }  // "Diffuse colour"
		file.float(tag: "Ks", count: 3) { f in mat!.specular = Colour(r: f[0], g: f[1], b: f[2]) }  // "Specular colour"
		file.float(tag: "Ns", count: 1) { f in mat!.specularExp = f[0] }  // "Specular exponent/shininess"
		file.float(tag: "Ni", count: 1) { f in mat!.refraction = f[0] }  // "Optical density/refraction index"
		file.float(tag: "d", count: 1) { f in mat!.alpha = f[0] }  // "Dissolve" 0.0 = transparent, 1.0 = opaque
		file.float(tag: "Tr", count: 1) { f in mat!.alpha = 1.0 - f[0] }  // "Transparent" 0.0 = opaque, 1.0 = transparent
		file.int(tag: "illum", count: 1) { i in mat!.model = try .resolve(i[0]) }  // "Illumination model"
		file.textureMap(tag: "map_Ka") { s in }  // "Ambient colourmap"
		file.textureMap(tag: "map_Kd") { t in mat!.diffuseMap = t }  // "Albedo map"
		file.textureMap(tag: "map_Ks") { t in mat!.specularMap = t }  // "Specular colourmap"
		file.textureMap(tag: "map_Ns") { t in mat!.specularExpMap = t }  // "Specular exponent map"
		file.textureMap(tag: "map_d") { t in mat!.dissolveMap = t }  // "Dissolve map"
		//file.string(tag: "map_Tr") { t in }  // "Translucency map"
		for t in ["map_Bump", "bump"] { file.textureMap(tag: t) { t in mat!.bumpMap = t } }  // "Bump map"
		file.textureMap(tag: "disp") { t in mat!.displacementMap = t }  // "Displacement map"
		file.textureMap(tag: "decal") { t in mat!.decalMap = t }  // ?
		// PBR extensions
		file.float(tag: "Pr", count: 1) { f in mat!.roughness = f[0] }  // "Roughness"
		file.float(tag: "Pm", count: 1) { f in mat!.metallic = f[0] }  // "Metallic"
		file.float(tag: "Ps", count: 1) { f in mat!.sheen = f[0] }  // "Sheen"
		file.float(tag: "Pc", count: 1) { f in }  // "Clearcoat thickness"
		file.float(tag: "Pcr", count: 1) { f in }  // "Clearcoat roughness"
		file.float(tag: "Ke", count: 3) { f in }  // "Emmission colour"
		file.float(tag: "aniso", count: 1) { f in }  // "Anisotropy"
		file.float(tag: "anisor", count: 1) { f in }  // "Anisotropy rotation"
		file.textureMap(tag: "map_Pr") { t in mat!.roughnessMap = t }  // "Roughness texturemap"
		file.textureMap(tag: "map_Pm") { t in mat!.metallicMap = t }  // "Metallic texturemap"
		file.textureMap(tag: "map_Ps") { t in mat!.sheenMap = t }  // "Sheen texturemap"
		file.textureMap(tag: "map_Ke") { t in mat!.emmisionMap = t }  // "Emmision texturemap"
		file.textureMap(tag: "norm") { t in mat!.normalMap = t }  // "Normal map"

		try file.read()

		if let material = mat
		{
			materials[name] = material
		}
		return materials
	}
}

fileprivate extension ObjMaterial.IlluminationModel
{
	static func resolve(_ id: Int) throws -> Self
	{
		switch id
		{
		case 0: .colour
		case 1: .lambert
		case 2: .blinnPhong
		case 3: .reflectionRaytrace
		case 4: .transparentGlassReflectionRaytrace
		case 5: .reflectionRaytraceFresnel
		case 6: .transparentRefractionReflectionRaytrace
		case 7: .transparentRefractionReflectionRaytraceFresnel
		case 8: .reflection
		case 9: .transparentGlassReflection
		case 10: .shadowOnly
		default: throw ObjLoaderError.badTagParameters
		}
	}
}

fileprivate extension ObjMaterial.TextureMap
{
	static func parse(_ argstr: Substring) throws -> Self
	{
		var map = Self()
		var args = [Substring]()

		let parseFlag =
		{ (flag: Flags) in
			let arg =  args[0].lowercased()
			if arg == "on" { map.flags.insert(flag) }
			else if arg == "off" { map.flags.remove(flag) }
			else { throw ObjLoaderError.badTagParameters }
		}
		let parseInt =
		{ (arg: Int) in
			guard let result = Int(args[arg]) else { throw ObjLoaderError.badTagParameters }
			return result
		}
		let parseFloat =
		{ (arg: Int) in
			guard let result = Float(args[arg]) else { throw ObjLoaderError.badTagParameters }
			return result
		}
		let parseVector = { Vec3f(try parseFloat(0), try parseFloat(1), try parseFloat(2)) }

		typealias Option = () throws -> Void
		let options: Dictionary<String, (Int, Option)> =
		[
			"blendu":  (1, { try parseFlag(.blendHoriz) }),
			"blendv":  (1, { try parseFlag(.blendVert) }),
			"bm":      (1, { map.blendMul = try parseFloat(0) }),
			"cc":      (1, { try parseFlag(.colourCorrection) }),
			"clamp":   (1, { try parseFlag(.clamp) }),
			"imfchan": (1, { map.imfChan = try .resolve(args[0]) }),
			"mm":      (2, { map.mmBaseGain = (try parseFloat(0), try parseFloat(1)) }),
			"o":       (3, { map.offset = try parseVector() }),
			"s":       (3, { map.scale = try parseVector() }),
			"t":       (3, { map.turbulence = try parseVector() }),
			"texres":  (1, { map.textureResolution = try parseInt(0) })
		]

		var expectArgs = 0, option: Option? = nil
		var index = argstr.startIndex
		repeat
		{
			if expectArgs > 0, let callback = option
			{
				let start = index
				index = argstr[start...].firstIndex(where: \.isWhitespace) ?? argstr.endIndex
				args.append(argstr[start..<index])
				expectArgs -= 1
				if expectArgs == 0
				{
					try callback()
					option = nil
				}
			}
			else if argstr[index] == "-"
			{
				let start = argstr.index(after: index)
				index = argstr[start...].firstIndex(where: \.isWhitespace) ?? argstr.endIndex
				let name = String(argstr[start..<index])
				guard let params = options[name] else { throw ObjLoaderError.badTagParameters }
				(expectArgs, option) = params
			}
			else
			{
				break
			}
			index = argstr[index...].firstIndex { !$0.isWhitespace } ?? argstr.endIndex
		}
		while index < argstr.endIndex
		if index >= argstr.endIndex || expectArgs > 0 { throw ObjLoaderError.badTagParameters }

		map.path = String(argstr[index...])
		return map
	}
}

fileprivate extension ObjMaterial.TextureMap.ImfChan
{
	static func resolve(_ token: Substring) throws -> Self
	{
		switch token
		{
		case "r": .r
		case "g": .g
		case "b": .b
		case "m": .m
		case "l": .l
		case "z": .z
		default: throw ObjLoaderError.badTagParameters
		}
	}
}

fileprivate struct ObjDocumentReader
{
	private enum Handler
	{
		case string(closure: ([String]) throws -> Void, count: Int? = nil)
		case float(closure: ([Float]) throws -> Void, count: Int? = nil)
		case int(closure: ([Int]) throws -> Void, count: Int? = nil)
		case textureMap(closure: (ObjMaterial.TextureMap) throws -> Void)
		case raw(closure: ([Substring]) throws -> Void)
	}

	private var file: TextFile
	typealias Callback = (Substring) throws -> Void
	private var _prehandler: Callback? = nil
	private var handlers = Dictionary<String, Handler>()

	init(filePath: URL) throws
	{
		file = try TextFile(fileURL: filePath)
	}

	mutating func string(tag: String, count: Int? = nil, closure: @escaping ([String]) throws -> Void)
	{
		handlers[tag] = .string(closure: closure, count: count)
	}

	mutating func float(tag: String, count: Int? = nil, closure: @escaping ([Float]) throws -> Void)
	{
		handlers[tag] = .float(closure: closure, count: count)
	}

	mutating func int(tag: String, count: Int? = nil, closure: @escaping ([Int]) throws -> Void)
	{
		handlers[tag] = .int(closure: closure, count: count)
	}

	mutating func textureMap(tag: String, closure: @escaping (ObjMaterial.TextureMap) throws -> Void)
	{
		handlers[tag] = .textureMap(closure: closure)
	}

	mutating func raw(tag: String, closure: @escaping ([Substring]) throws -> Void)
	{
		handlers[tag] = .raw(closure: closure)
	}

	mutating func preHandle(closure: @escaping Callback) { _prehandler = closure }

	private func handle(_ handler: Handler, command: Substring, arg: Substring) throws
	{
		switch handler
		{
		case .string(let closure, let count):
			let args = arg.split(separator: " ")
			if count != nil && args.count != count! { throw ObjLoaderError.badTagParameters }
			try closure(args.map({ String($0) }))
		case .float(let closure, let count):
			let args = arg.split(separator: " ")
			if count != nil && args.count != count! { throw ObjLoaderError.badTagParameters }
			try closure(args.map(
			{
				if let value = Float($0) { return value }
				else { throw ObjLoaderError.badTagParameters }
			}))
		case .int(let closure, let count):
			let args = arg.split(separator: " ")
			if count != nil && args.count != count! { throw ObjLoaderError.badTagParameters }
			try closure(args.map(
			{
				if let value = Int($0) { return value }
				else { throw ObjLoaderError.badTagParameters }
			}))
		case .textureMap(let closure):
			try closure(ObjMaterial.TextureMap.parse(arg))
		case .raw(let closure):
			try closure(arg.split(separator: " "))
		}
	}

	private func parseLine(_ line: String) throws -> (Substring, Substring)?
	{
		// Trim comment if present
		var trimmed = if let index = line.firstIndex(of: "#") { line[..<index] } else { line[...] }

		// Trim leading and trailing whitespace
		guard let start = trimmed.firstIndex(where: { !$0.isWhitespace }),
			let end = trimmed.lastIndex(where: { !$0.isWhitespace })
		else { return nil }
		trimmed = trimmed[start...end]

		// Split command & rest of string as arguments
		guard let cmdEnd = trimmed.firstIndex(where: \.isWhitespace),
			let argStart = trimmed.suffix(from: trimmed.index(after: cmdEnd)).firstIndex(where: { !$0.isWhitespace })
		else { throw ObjLoaderError.badTagParameters }
		return (trimmed.prefix(upTo: cmdEnd), trimmed.suffix(from: argStart))
	}

	func read() throws
	{
		for line in file.lines
		{
			if !line.isEmpty,
				let (cmd, argstr) = try parseLine(line),
				let handler = handlers[String(cmd)]
			{
				if let prehandler = _prehandler { try prehandler(cmd) }
				try handle(handler, command: cmd, arg: argstr)
			}
		}
	}
}


enum ObjLoaderError: Error
{
	case badTagParameters
	case unexpectedTag
}