import SwiftUI
import AppKit

private struct DiagramRelation: Identifiable {
    var id: String { "\(from).\(column).\(to).\(targetColumn)" }
    let from: String
    let column: String
    let to: String
    let targetColumn: String
}

private struct DiagramPoint: Codable {
    var x: Double
    var y: Double
}

// Keep the drag translation local to the node. Updating the entire diagram's
// positions dictionary on every mouse event caused layout churn and jerky drags.
private struct DiagramDraggableNode<Content: View>: View {
    let zoom: CGFloat
    let onDrop: (CGSize) -> Void
    @ViewBuilder let content: () -> Content
    @GestureState private var translation: CGSize = .zero

    var body: some View {
        content()
            .offset(x: translation.width / max(zoom, 0.1),
                    y: translation.height / max(zoom, 0.1))
            .highPriorityGesture(
                DragGesture(minimumDistance: 3, coordinateSpace: .global)
                    .updating($translation) { value, state, _ in
                        state = value.translation
                    }
                    .onEnded { value in
                        onDrop(value.translation)
                    }
            )
    }
}

@MainActor
struct DBSynSchemaDiagramView: View {
    @ObservedObject var model: DatabaseModel
    let database: String
    var onClose: () -> Void = {}
    @State private var positions: [String: CGPoint] = [:]
    @State private var relations: [DiagramRelation] = []
    @State private var zoom: CGFloat = 1
    @State private var search = ""
    @State private var selectedObject: String?
    @State private var loadingRelationships = false
    @State private var errorMessage: String?
    @State private var showViews = true
    @State private var showingObjects = true
    @State private var includedObjects: Set<String> = []
    @State private var selectionRevision = 0
    @State private var theme = "Ocean"
    @State private var showAllColumns = false
    @State private var relatedOnly = false
    @State private var diagramName = "Main"
    @State private var diagrams: [String] = ["Main"]
    @State private var newDiagramName = ""
    @State private var showingNewDiagram = false
    @State private var showingDeleteConfirmation = false
    @State private var exportError: String?
    @State private var showingExportError = false

    private let nodeWidth: CGFloat = 245
    private let headerHeight: CGFloat = 36
    private let rowHeight: CGFloat = 21
    private var maxColumns: Int { showAllColumns ? 200 : 10 }
    private var snapshot: DatabaseModel.SchemaIntrospection? { model.introspection(for: database) }
    private var objects: [String] {
        guard let snapshot else { return [] }
        let all = snapshot.tables + (showViews ? snapshot.views : [])
        return all.filter { name in
            includedObjects.contains(name) && (search.isEmpty || name.localizedCaseInsensitiveContains(search)) &&
            (!relatedOnly || selectedObject == nil || name == selectedObject || relations.contains { ($0.from == selectedObject && $0.to == name) || ($0.to == selectedObject && $0.from == name) })
        }
    }
    private var diagramNamespace: String {
        "DBSyn.diagram.v2.\(model.selected?.id.uuidString ?? "unknown").\(database)"
    }
    private var storageKey: String { "\(diagramNamespace).layout.\(diagramName)" }
    private var catalogKey: String { "\(diagramNamespace).catalog" }
    private var selectionKey: String { "\(diagramNamespace).objects.\(diagramName)" }
    private var themeKey: String { "\(diagramNamespace).theme.\(diagramName)" }
    private var allObjects: [String] { (snapshot?.tables ?? []) + (snapshot?.views ?? []) }
    private var nodeColor: Color {
        switch theme {
        case "Slate": return Color(red: 0.30, green: 0.38, blue: 0.49)
        case "Forest": return Color(red: 0.16, green: 0.43, blue: 0.36)
        case "Plum": return Color(red: 0.42, green: 0.32, blue: 0.55)
        default: return Color(red: 0.12, green: 0.38, blue: 0.66)
        }
    }
    private func restoreOptions() {
        if let stored = UserDefaults.standard.stringArray(forKey: selectionKey) {
            includedObjects = Set(stored)
        } else { includedObjects = Set(allObjects) }
        theme = UserDefaults.standard.string(forKey: themeKey) ?? UserDefaults.standard.string(forKey: "preferences.diagramPalette") ?? "Ocean"
        showAllColumns = UserDefaults.standard.bool(forKey: "preferences.diagramShowAllColumns")
    }
    private func saveOptions() {
        UserDefaults.standard.set(Array(includedObjects).sorted(), forKey: selectionKey)
        UserDefaults.standard.set(theme, forKey: themeKey)
    }
    private func columns(_ name: String) -> [DatabaseModel.IntrospectedColumn] {
        snapshot?.columnsByTable[name] ?? []
    }
    private func height(_ name: String) -> CGFloat {
        headerHeight + CGFloat(min(columns(name).count, maxColumns)) * rowHeight + 12
    }
    private func position(_ name: String) -> CGPoint {
        positions[name] ?? CGPoint(x: 30, y: 30)
    }
    private func center(_ name: String, right: Bool) -> CGPoint {
        let p = position(name)
        return CGPoint(x: p.x + (right ? nodeWidth : 0), y: p.y + min(height(name), 260) / 2)
    }

    var body: some View {
        VStack(spacing: 0) {
            // Two compact toolbar rows avoid the truncated controls shown in the original sheet.
            HStack(spacing: 10) {
                Label(database, systemImage: "point.3.connected.trianglepath.dotted")
                    .font(.headline).lineLimit(1)
                Text("SCHEMA DESIGNER").font(.caption2.weight(.semibold)).foregroundStyle(.secondary)
                Spacer(minLength: 12)
                Button("Close", systemImage: "xmark") { savePositions(); saveOptions(); onClose() }
                    .labelStyle(.iconOnly).help("Close diagram")
                    .keyboardShortcut(.cancelAction)
            }
            .padding(.horizontal, 16).padding(.top, 12).padding(.bottom, 8)
            HStack(spacing: 10) {
                Picker("Diagram", selection: $diagramName) {
                    ForEach(diagrams, id: \.self) { Text($0).tag($0) }
                }
                .frame(width: 175)
                .onChange(of: diagramName) { oldValue, _ in
                    savePositions(named: oldValue)
                    positions = [:]
                    restorePositions()
                    restoreOptions()
                    if positions.isEmpty { arrange() }
                }
                Button { showingNewDiagram = true } label: { Image(systemName: "plus") }
                    .help("New diagram")
                Button { showingDeleteConfirmation = true } label: { Image(systemName: "trash") }
                    .disabled(diagrams.count <= 1).help("Delete diagram")
                Divider().frame(height: 18)
                Button { showingObjects.toggle() } label: {
                    Label("Tables", systemImage: "sidebar.left")
                }.help("Show or hide the table selection sidebar")
                TextField("Find table or view", text: $search)
                    .textFieldStyle(.roundedBorder).frame(minWidth: 130, idealWidth: 220, maxWidth: 260)
                Toggle("Views", isOn: $showViews).toggleStyle(.checkbox)
                Toggle("Related", isOn: $relatedOnly).toggleStyle(.checkbox)
                Spacer(minLength: 8)
                Menu {
                    ForEach(["Ocean", "Slate", "Forest", "Plum"], id: \.self) { option in
                        Button { theme = option; saveOptions() } label: {
                            if theme == option { Label(option, systemImage: "checkmark") }
                            else { Text(option) }
                        }
                    }
                } label: { Label("Colours", systemImage: "paintpalette") }
                Menu {
                    Toggle("Show all columns", isOn: $showAllColumns)
                    Button("Auto Layout") { arrange(); savePositions() }
                    Button("Refresh Metadata") { model.introspect(database: database); loadRelationships() }
                } label: { Label("Layout", systemImage: "slider.horizontal.3") }
                Button { zoom = max(0.4, zoom - 0.15) } label: { Image(systemName: "minus.magnifyingglass") }
                Text("\(Int(zoom * 100))%").monospacedDigit().frame(width: 45)
                Button { zoom = min(2, zoom + 0.15) } label: { Image(systemName: "plus.magnifyingglass") }
                Menu {
                    Button("Export PDF…") { exportDiagram(.pdf) }
                    Button("Export PNG…") { exportDiagram(.png) }
                    Button("Export SVG…") { exportDiagram(.svg) }
                    Divider()
                    Button("Print…") { printDiagram() }
                } label: { Label("Export / Print", systemImage: "square.and.arrow.up") }
            }
            .controlSize(.small)
            .padding(.horizontal, 16).padding(.bottom, 12)
            Divider()
            HStack(spacing: 0) {
                if showingObjects {
                    VStack(alignment: .leading, spacing: 10) {
                        HStack {
                            Text("OBJECTS").font(.caption.weight(.semibold)).foregroundStyle(.secondary)
                            Spacer()
                            Text("\(includedObjects.count)/\(allObjects.count)")
                                .font(.caption.monospacedDigit()).foregroundStyle(.secondary)
                        }
                        HStack {
                            Button("All") { updateIncludedObjects(Set(allObjects)) }
                            Button("None") { updateIncludedObjects([]) }
                            Spacer()
                        }.buttonStyle(.link)
                        ScrollView {
                            LazyVStack(alignment: .leading, spacing: 5) {
                                ForEach(allObjects.filter { search.isEmpty || $0.localizedCaseInsensitiveContains(search) }, id: \.self) { name in
                                    Toggle(isOn: Binding(
                                        get: { includedObjects.contains(name) },
                                        set: { enabled in
                                            var updated = includedObjects
                                            if enabled { updated.insert(name) }
                                            else { updated.remove(name) }
                                            updateIncludedObjects(updated)
                                        })) {
                                        Label(name, systemImage: snapshot?.views.contains(name) == true ? "eye" : "tablecells")
                                            .lineLimit(1).font(.system(size: 12))
                                    }.toggleStyle(.checkbox)
                                }
                            }.frame(maxWidth: .infinity, alignment: .leading)
                        }
                    }
                    .padding(12).frame(width: 235)
                    .background(Color(nsColor: .controlBackgroundColor))
                    Divider()
                }
                VStack(spacing: 0) {
            if snapshot == nil {
                ContentUnavailableView("No Introspection Data", systemImage: "externaldrive.badge.questionmark", description: Text("Run Introspection on this database to load tables and views."))
                Button("Introspect Now") { model.introspect(database: database); loadRelationships() }.padding()
            } else {
                ScrollView([.horizontal, .vertical]) {
                    ZStack(alignment: .topLeading) {
                        Canvas { context, _ in
                            for relation in relations where objects.contains(relation.from) && objects.contains(relation.to) {
                                let a = position(relation.from)
                                let b = position(relation.to)
                                let right = a.x <= b.x
                                let start = center(relation.from, right: right)
                                let end = center(relation.to, right: !right)
                                let bend: CGFloat = right ? 65 : -65
                                var path = Path()
                                path.move(to: start)
                                path.addCurve(to: end,
                                              control1: CGPoint(x: start.x + bend, y: start.y),
                                              control2: CGPoint(x: end.x - bend, y: end.y))
                                context.stroke(path, with: .color(.accentColor.opacity(0.7)), lineWidth: 1.7)
                            }
                        }
                        ForEach(objects, id: \.self) { name in
                            DiagramDraggableNode(zoom: zoom, onDrop: { translation in
                                let start = position(name)
                                positions[name] = CGPoint(
                                    x: max(10, start.x + translation.width / max(zoom, 0.1)),
                                    y: max(10, start.y + translation.height / max(zoom, 0.1)))
                                savePositions()
                            }) {
                                diagramNode(name)
                            }
                            .position(x: position(name).x + nodeWidth / 2,
                                      y: position(name).y + height(name) / 2)
                        }
                    }
                    .frame(width: canvasWidth, height: canvasHeight, alignment: .topLeading)
                    .scaleEffect(zoom, anchor: .topLeading)
                    .frame(width: canvasWidth * zoom, height: canvasHeight * zoom, alignment: .topLeading)
                }
                .background(Color(nsColor: .textBackgroundColor))
                .id(selectionRevision)
                .onChange(of: snapshot?.loadedAt) { _, _ in
                    if UserDefaults.standard.stringArray(forKey: selectionKey) == nil {
                        includedObjects = Set(allObjects)
                    }
                    if positions.isEmpty { arrange() }
                    loadRelationships()
                }
            }
                }
            }
            Divider()
            HStack {
                Text("\(snapshot?.tables.count ?? 0) tables · \(snapshot?.views.count ?? 0) views · \(relations.count) foreign-key relationships")
                if loadingRelationships { ProgressView().controlSize(.small) }
                if let errorMessage { Text(errorMessage).foregroundStyle(.secondary).lineLimit(1) }
                Spacer()
                Text("Drag tables to arrange · Select to highlight relationships")
            }
            .font(.caption)
            .padding(9)
        }
        .frame(minWidth: 850, idealWidth: 1320, minHeight: 570, idealHeight: 850)
        .background(DiagramResizableWindow())
        .onAppear { loadCatalog(); restorePositions(); restoreOptions(); if positions.isEmpty { arrange() }; loadRelationships() }
        .onDisappear { savePositions(); saveOptions() }
        .alert("New Schema Diagram", isPresented: $showingNewDiagram) {
            TextField("Diagram name", text: $newDiagramName)
            Button("Create") { createDiagram() }
            Button("Cancel", role: .cancel) { }
        } message: { Text("Create a separate saved layout for this database.") }
        .confirmationDialog("Delete \(diagramName)?", isPresented: $showingDeleteConfirmation) {
            Button("Delete Diagram", role: .destructive) { deleteDiagram() }
        } message: { Text("Only the saved diagram layout is removed. Database objects are not changed.") }
        .alert("Diagram Export", isPresented: $showingExportError) {
            Button("OK", role: .cancel) { }
        } message: { Text(exportError ?? "Unable to export diagram") }
    }

    private var canvasWidth: CGFloat {
        max(1200, objects.map { position($0).x + nodeWidth + 80 }.max() ?? 1200)
    }
    private var canvasHeight: CGFloat {
        max(800, objects.map { position($0).y + height($0) + 80 }.max() ?? 800)
    }

    private func diagramNode(_ name: String) -> some View {
        let isView = snapshot?.views.contains(name) == true
        let related = selectedObject == nil || selectedObject == name || relations.contains { ($0.from == selectedObject && $0.to == name) || ($0.to == selectedObject && $0.from == name) }
        return VStack(alignment: .leading, spacing: 0) {
            HStack(spacing: 7) {
                Image(systemName: isView ? "eye" : "tablecells")
                Text(name).fontWeight(.semibold).lineLimit(1)
                Spacer()
                if isView { Text("VIEW").font(.system(size: 9, weight: .bold)) }
            }
            .font(.system(size: 12))
            .padding(.horizontal, 10)
            .frame(height: headerHeight)
            .foregroundStyle(.white)
            .background(isView ? nodeColor.opacity(0.76) : nodeColor)
            ForEach(Array(columns(name).prefix(maxColumns)), id: \.id) { column in
                HStack(spacing: 5) {
                    Image(systemName: column.key == "PRI" ? "key.fill" : "circle.fill")
                        .font(.system(size: column.key == "PRI" ? 10 : 4))
                        .foregroundStyle(column.key == "PRI" ? Color.orange : Color.secondary)
                        .frame(width: 13)
                    Text(column.name).lineLimit(1)
                    Spacer(minLength: 4)
                    Text(column.dataType).foregroundStyle(.secondary).lineLimit(1)
                }
                .font(.system(size: 10))
                .padding(.horizontal, 9)
                .frame(height: rowHeight)
            }
            if columns(name).count > maxColumns {
                Text("+ \(columns(name).count - maxColumns) more columns")
                    .font(.caption2).foregroundStyle(.secondary).padding(.horizontal, 10)
            }
            Spacer(minLength: 5)
        }
        .frame(width: nodeWidth, height: height(name))
        .background(Color(nsColor: .controlBackgroundColor))
        .clipShape(RoundedRectangle(cornerRadius: 7))
        .overlay(RoundedRectangle(cornerRadius: 7).stroke(selectedObject == name ? Color.accentColor : Color(nsColor: .separatorColor), lineWidth: selectedObject == name ? 2.5 : 1))
        .shadow(color: .black.opacity(0.13), radius: 3, y: 2)
        .opacity(related ? 1 : 0.4)
        .onTapGesture { selectedObject = selectedObject == name ? nil : name }
    }

    // Rebuild the visible canvas whenever the selected objects change. The
    // previous implementation only saved the checkboxes, leaving the diagram
    // scrolled to stale positions after None / All / individual selections.
    private func updateIncludedObjects(_ updated: Set<String>) {
        let valid = DiagramLogic.selectedObjects(updated, available: allObjects)
        guard valid != includedObjects else { return }
        includedObjects = valid
        if let selectedObject, !valid.contains(selectedObject) {
            self.selectedObject = nil
        }
        // Keep the selected subset compact, rather than retaining the old
        // coordinates of objects from a much larger diagram.
        arrange()
        savePositions()
        saveOptions()
        selectionRevision += 1
    }

    private func arrange() {
        let names = allObjects.filter { includedObjects.contains($0) }
        positions.merge(DiagramLogic.arrangedPositions(names)) { _, latest in latest }
    }
    private func loadCatalog() {
        let stored = UserDefaults.standard.stringArray(forKey: catalogKey) ?? ["Main"]
        diagrams = stored.isEmpty ? ["Main"] : stored
        if !diagrams.contains(diagramName) { diagramName = diagrams[0] }
    }
    private func createDiagram() {
        let name = newDiagramName.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !name.isEmpty, !diagrams.contains(name) else { return }
        savePositions()
        diagrams.append(name)
        UserDefaults.standard.set(diagrams, forKey: catalogKey)
        diagramName = name
        newDiagramName = ""
    }
    private func deleteDiagram() {
        guard diagrams.count > 1 else { return }
        UserDefaults.standard.removeObject(forKey: storageKey)
        diagrams.removeAll { $0 == diagramName }
        UserDefaults.standard.set(diagrams, forKey: catalogKey)
        diagramName = diagrams[0]
    }
    private func savePositions(named name: String) {
        let points = positions.mapValues { DiagramPoint(x: Double($0.x), y: Double($0.y)) }
        if let data = try? JSONEncoder().encode(points) {
            UserDefaults.standard.set(data, forKey: "\(diagramNamespace).layout.\(name)")
        }
    }
    private func savePositions() { savePositions(named: diagramName) }
    private func restorePositions() {
        guard let data = UserDefaults.standard.data(forKey: storageKey),
              let saved = try? JSONDecoder().decode([String: DiagramPoint].self, from: data) else { return }
        positions = saved.mapValues { CGPoint(x: $0.x, y: $0.y) }
    }
    private func loadRelationships() {
        guard let profile = model.selected else { return }
        loadingRelationships = true
        errorMessage = nil
        let password = PasswordStore.read(profile.id)
        let escaped = database.replacingOccurrences(of: "'", with: "''")
        Task {
            do {
                let result = try await NativeMySQLRunner.run(profile: profile, database: database, password: password,
                    sql: "SELECT TABLE_NAME, COLUMN_NAME, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME FROM information_schema.KEY_COLUMN_USAGE WHERE TABLE_SCHEMA = '\(escaped)' AND REFERENCED_TABLE_NAME IS NOT NULL ORDER BY TABLE_NAME, ORDINAL_POSITION")
                guard model.selected?.id == profile.id else { loadingRelationships = false; return }
                relations = result.rows.compactMap { row in
                    guard row.count >= 4 else { return nil }
                    return DiagramRelation(from: row[0], column: row[1], to: row[2], targetColumn: row[3])
                }
            } catch {
                errorMessage = "Foreign keys unavailable: \(error.localizedDescription)"
            }
            loadingRelationships = false
        }
    }
    private enum DiagramFormat: Equatable { case pdf, png, svg }

    private func printLayout() -> SchemaDiagramPrintView {
        let entries = objects.map { name in
            SchemaDiagramPrintView.Entry(name: name,
                isView: snapshot?.views.contains(name) == true,
                columns: columns(name).map { .init(name: $0.name, type: $0.dataType, primary: $0.key == "PRI") },
                origin: position(name))
        }
        let links = relations.filter { objects.contains($0.from) && objects.contains($0.to) }
            .map { SchemaDiagramPrintView.Link(from: $0.from, to: $0.to, column: $0.column, target: $0.targetColumn) }
        return SchemaDiagramPrintView(database: database, diagram: diagramName, entries: entries, links: links, theme: theme)
    }

    private func exportDiagram(_ format: DiagramFormat) {
        let panel = NSSavePanel()
        panel.title = "Export Schema Diagram"
        panel.nameFieldStringValue = "\(database)-\(diagramName).\(format == .pdf ? "pdf" : format == .png ? "png" : "svg")"
        panel.canCreateDirectories = true
        guard panel.runModal() == .OK, let url = panel.url else { return }
        do {
            let view = printLayout()
            let data: Data
            switch format {
            case .pdf: data = view.dataWithPDF(inside: view.bounds)
            case .png:
                guard let bitmap = view.bitmapImageRepForCachingDisplay(in: view.bounds) else {
                    throw NSError(domain: "DBSynDiagram", code: 1, userInfo: [NSLocalizedDescriptionKey: "Unable to create bitmap"])
                }
                view.cacheDisplay(in: view.bounds, to: bitmap)
                guard let png = bitmap.representation(using: .png, properties: [:]) else {
                    throw NSError(domain: "DBSynDiagram", code: 2, userInfo: [NSLocalizedDescriptionKey: "Unable to encode PNG"])
                }
                data = png
            case .svg: data = Data(view.svgDocument().utf8)
            }
            try data.write(to: url, options: .atomic)
        } catch {
            exportError = error.localizedDescription
            showingExportError = true
        }
    }

    private func printDiagram() {
        let view = printLayout()
        guard let info = NSPrintInfo.shared.copy() as? NSPrintInfo else { return }
        info.horizontalPagination = .fit
        info.verticalPagination = .fit
        info.isHorizontallyCentered = true
        info.isVerticallyCentered = true
        info.orientation = view.bounds.width > view.bounds.height ? .landscape : .portrait
        NSPrintOperation(view: view, printInfo: info).run()
    }

}

/// Allows the schema designer sheet to be resized rather than locking its initial dimensions.
private struct DiagramResizableWindow: NSViewRepresentable {
    func makeNSView(context: Context) -> NSView {
        let view = NSView(frame: .zero)
        DispatchQueue.main.async { [weak view] in
            guard let window = view?.window else { return }
            window.styleMask.insert(.resizable)
            window.minSize = NSSize(width: 850, height: 570)
        }
        return view
    }
    func updateNSView(_ nsView: NSView, context: Context) { }
}