import Foundation

public struct ComparisonResult {
    public let audioResults: [AudioTestResult]
    public let visionResults: [VisionTestResult]
    public let endToEndResults: [EndToEndResult]
    public let timestamp: Date
}

public struct EndToEndResult {
    public let modelName: String
    public let loadTimeMs: Double
    public let audioProcessMs: Double
    public let visionProcessMs: Double
    public let genSpeedTokPerSec: Double
    public let totalTimeMs: Double
    public let generatedTokens: Int
    public let passed: Bool
}

public final class ComparisonReporter {
    private let outputDir: String
    
    public init(outputDir: String = "/tmp/multimodal_comparison") {
        self.outputDir = outputDir
        try? FileManager.default.createDirectory(atPath: outputDir, withIntermediateDirectories: true)
    }
    
    public func printAudioTable(results: [AudioTestResult], sampleName: String) {
        print("\n▶ AUDIO TEST: \(sampleName)")
        printTableHeader(columns: ["Metric", "E2B", "E4B", "12B"], widths: [12, 15, 15, 15])
        printTableSeparator(widths: [12, 15, 15, 15])
        
        let e2b = results.first { $0.modelName == "E2B" }
        let e4b = results.first { $0.modelName == "E4B" }
        let g12b = results.first { $0.modelName == "12B" }
        
        printTableRow(columns: ["Shape", 
            e2b != nil ? "[\(e2b!.outputShape.0), \(e2b!.outputShape.1)]" : "N/A",
            e4b != nil ? "[\(e4b!.outputShape.0), \(e4b!.outputShape.1)]" : "N/A",
            g12b != nil ? "[\(g12b!.outputShape.0), \(g12b!.outputShape.1)]" : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Min",
            e2b != nil ? String(format: "%.3f", e2b!.min) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.min) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.min) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Max",
            e2b != nil ? String(format: "%.3f", e2b!.max) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.max) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.max) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Mean",
            e2b != nil ? String(format: "%.3f", e2b!.mean) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.mean) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.mean) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Std",
            e2b != nil ? String(format: "%.3f", e2b!.std) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.std) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.std) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["NaN",
            e2b != nil ? String(e2b!.nanCount) : "N/A",
            e4b != nil ? String(e4b!.nanCount) : "N/A",
            g12b != nil ? String(g12b!.nanCount) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Time",
            e2b != nil ? String(format: "%.1f ms", e2b!.forwardTimeMs) : "N/A",
            e4b != nil ? String(format: "%.1f ms", e4b!.forwardTimeMs) : "N/A",
            g12b != nil ? String(format: "%.1f ms", g12b!.forwardTimeMs) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Memory",
            e2b != nil ? String(format: "%.0f MB", e2b!.memoryPeakMB) : "N/A",
            e4b != nil ? String(format: "%.0f MB", e4b!.memoryPeakMB) : "N/A",
            g12b != nil ? String(format: "%.0f MB", g12b!.memoryPeakMB) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Status",
            e2b != nil ? (e2b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A",
            e4b != nil ? (e4b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A",
            g12b != nil ? (g12b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableSeparator(widths: [12, 15, 15, 15])
    }
    
    public func printVisionTable(results: [VisionTestResult], sampleName: String) {
        print("\n▶ VISION TEST: \(sampleName)")
        printTableHeader(columns: ["Metric", "E2B", "E4B", "12B"], widths: [12, 15, 15, 15])
        printTableSeparator(widths: [12, 15, 15, 15])
        
        let e2b = results.first { $0.modelName == "E2B" }
        let e4b = results.first { $0.modelName == "E4B" }
        let g12b = results.first { $0.modelName == "12B" }
        
        printTableRow(columns: ["Shape",
            e2b != nil ? "[\(e2b!.outputShape.0), \(e2b!.outputShape.1)]" : "N/A",
            e4b != nil ? "[\(e4b!.outputShape.0), \(e4b!.outputShape.1)]" : "N/A",
            g12b != nil ? "[\(g12b!.outputShape.0), \(g12b!.outputShape.1)]" : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Min",
            e2b != nil ? String(format: "%.3f", e2b!.min) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.min) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.min) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Max",
            e2b != nil ? String(format: "%.3f", e2b!.max) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.max) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.max) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Mean",
            e2b != nil ? String(format: "%.3f", e2b!.mean) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.mean) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.mean) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Std",
            e2b != nil ? String(format: "%.3f", e2b!.std) : "N/A",
            e4b != nil ? String(format: "%.3f", e4b!.std) : "N/A",
            g12b != nil ? String(format: "%.3f", g12b!.std) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["NaN",
            e2b != nil ? String(e2b!.nanCount) : "N/A",
            e4b != nil ? String(e4b!.nanCount) : "N/A",
            g12b != nil ? String(g12b!.nanCount) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["CosSim",
            "—",
            e4b != nil && e2b != nil ? String(format: "%.3f", e4b!.cosineSimilarity ?? 0) : "N/A",
            "—"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Time",
            e2b != nil ? String(format: "%.1f ms", e2b!.forwardTimeMs) : "N/A",
            e4b != nil ? String(format: "%.1f ms", e4b!.forwardTimeMs) : "N/A",
            g12b != nil ? String(format: "%.1f ms", g12b!.forwardTimeMs) : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableRow(columns: ["Status",
            e2b != nil ? (e2b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A",
            e4b != nil ? (e4b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A",
            g12b != nil ? (g12b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A"
        ], widths: [12, 15, 15, 15])
        
        printTableSeparator(widths: [12, 15, 15, 15])
    }
    
    public func printEndToEndTable(results: [EndToEndResult]) {
        print("\n▶ END-TO-END TEST: audio+vision generation")
        printTableHeader(columns: ["Metric", "E2B", "E4B", "12B"], widths: [14, 12, 12, 12])
        printTableSeparator(widths: [14, 12, 12, 12])
        
        let e2b = results.first { $0.modelName == "E2B" }
        let e4b = results.first { $0.modelName == "E4B" }
        let g12b = results.first { $0.modelName == "12B" }
        
        printTableRow(columns: ["Load Time",
            e2b != nil ? String(format: "%.2f s", e2b!.loadTimeMs / 1000.0) : "N/A",
            e4b != nil ? String(format: "%.2f s", e4b!.loadTimeMs / 1000.0) : "N/A",
            g12b != nil ? String(format: "%.2f s", g12b!.loadTimeMs / 1000.0) : "N/A"
        ], widths: [14, 12, 12, 12])
        
        printTableRow(columns: ["Audio Proc",
            e2b != nil ? String(format: "%.1f ms", e2b!.audioProcessMs) : "N/A",
            e4b != nil ? String(format: "%.1f ms", e4b!.audioProcessMs) : "N/A",
            g12b != nil ? String(format: "%.1f ms", g12b!.audioProcessMs) : "N/A"
        ], widths: [14, 12, 12, 12])
        
        printTableRow(columns: ["Vision Proc",
            e2b != nil ? String(format: "%.1f ms", e2b!.visionProcessMs) : "N/A",
            e4b != nil ? String(format: "%.1f ms", e4b!.visionProcessMs) : "N/A",
            g12b != nil ? String(format: "%.1f ms", g12b!.visionProcessMs) : "N/A"
        ], widths: [14, 12, 12, 12])
        
        printTableRow(columns: ["Gen Speed",
            e2b != nil ? String(format: "%.1f t/s", e2b!.genSpeedTokPerSec) : "N/A",
            e4b != nil ? String(format: "%.1f t/s", e4b!.genSpeedTokPerSec) : "N/A",
            g12b != nil ? String(format: "%.1f t/s", g12b!.genSpeedTokPerSec) : "N/A"
        ], widths: [14, 12, 12, 12])
        
        printTableRow(columns: ["Total Time",
            e2b != nil ? String(format: "%.2f s", e2b!.totalTimeMs / 1000.0) : "N/A",
            e4b != nil ? String(format: "%.2f s", e4b!.totalTimeMs / 1000.0) : "N/A",
            g12b != nil ? String(format: "%.2f s", g12b!.totalTimeMs / 1000.0) : "N/A"
        ], widths: [14, 12, 12, 12])
        
        printTableRow(columns: ["Status",
            e2b != nil ? (e2b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A",
            e4b != nil ? (e4b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A",
            g12b != nil ? (g12b!.passed ? "✓ PASS" : "✗ FAIL") : "N/A"
        ], widths: [14, 12, 12, 12])
        
        printTableSeparator(widths: [14, 12, 12, 12])
    }
    
    public func printSummary() {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("                         SUMMARY")
        print("═══════════════════════════════════════════════════════════════════")
        print("Audio Tower:   E2B/E4B 完整12層，12B僅投影")
        print("Vision Tower:  E2B/E4B 完整16層，12B簡化版")
        print("Speed:         12B最快，但品質較差")
        print("Quality:       E4B最佳（量化精度平衡）")
        print("Memory:        E2B最小（2B參數）")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    private func printTableHeader(columns: [String], widths: [Int]) {
        var line = "┌"
        for (i, w) in widths.enumerated() {
            line += String(repeating: "─", count: w)
            if i < widths.count - 1 { line += "┬" }
        }
        line += "┐"
        print(line)
        
        var row = "│"
        for i in 0..<columns.count {
            let col = columns[i]
            let w = widths[i]
            row += padRight(col, width: w)
            if i < widths.count - 1 { row += "│" }
        }
        row += "│"
        print(row)
    }
    
    private func printTableSeparator(widths: [Int]) {
        var line = "├"
        for (i, w) in widths.enumerated() {
            line += String(repeating: "─", count: w)
            if i < widths.count - 1 { line += "┼" }
        }
        line += "┤"
        print(line)
    }
    
    private func printTableRow(columns: [String], widths: [Int]) {
        var row = "│"
        for i in 0..<columns.count {
            let col = columns[i]
            let w = widths[i]
            row += padRight(col, width: w)
            if i < widths.count - 1 { row += "│" }
        }
        row += "│"
        print(row)
    }
    
    private func padRight(_ s: String, width: Int) -> String {
        if s.count >= width { return s }
        return s + String(repeating: " ", count: width - s.count)
    }
    
    public func writeJSON(result: ComparisonResult) throws {
        let encoder = JSONEncoder()
        encoder.outputFormatting = .prettyPrinted
        let data = try encoder.encode(JSONReport(from: result))
        let path = outputDir + "/comparison_results.json"
        FileManager.default.createFile(atPath: path, contents: data)
        print("  JSON report saved to: \(path)")
    }
    
    public func writeCSV(result: ComparisonResult) throws {
        var csv = "type,sample,model,metric,value\n"
        
        for r in result.audioResults {
            csv += "audio,\(r.sampleName),\(r.modelName),min,\(r.min)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),max,\(r.max)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),mean,\(r.mean)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),std,\(r.std)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),nan_count,\(r.nanCount)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),time_ms,\(r.forwardTimeMs)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),memory_mb,\(r.memoryPeakMB)\n"
            csv += "audio,\(r.sampleName),\(r.modelName),passed,\(r.passed)\n"
        }
        
        for r in result.visionResults {
            csv += "vision,\(r.sampleName),\(r.modelName),min,\(r.min)\n"
            csv += "vision,\(r.sampleName),\(r.modelName),max,\(r.max)\n"
            csv += "vision,\(r.sampleName),\(r.modelName),mean,\(r.mean)\n"
            csv += "vision,\(r.sampleName),\(r.modelName),std,\(r.std)\n"
            csv += "vision,\(r.sampleName),\(r.modelName),nan_count,\(r.nanCount)\n"
            if let cs = r.cosineSimilarity {
                csv += "vision,\(r.sampleName),\(r.modelName),cosine_sim,\(cs)\n"
            }
            csv += "vision,\(r.sampleName),\(r.modelName),time_ms,\(r.forwardTimeMs)\n"
            csv += "vision,\(r.sampleName),\(r.modelName),passed,\(r.passed)\n"
        }
        
        for r in result.endToEndResults {
            csv += "e2e,all,\(r.modelName),load_time_ms,\(r.loadTimeMs)\n"
            csv += "e2e,all,\(r.modelName),audio_proc_ms,\(r.audioProcessMs)\n"
            csv += "e2e,all,\(r.modelName),vision_proc_ms,\(r.visionProcessMs)\n"
            csv += "e2e,all,\(r.modelName),gen_speed_tok_s,\(r.genSpeedTokPerSec)\n"
            csv += "e2e,all,\(r.modelName),total_time_ms,\(r.totalTimeMs)\n"
            csv += "e2e,all,\(r.modelName),passed,\(r.passed)\n"
        }
        
        let path = outputDir + "/comparison_results.csv"
        FileManager.default.createFile(atPath: path, contents: csv.data(using: .utf8))
        print("  CSV report saved to: \(path)")
    }
}

private struct JSONReport: Codable {
    let timestamp: String
    let audio: [AudioResultJSON]
    let vision: [VisionResultJSON]
    let endToEnd: [EndToEndJSON]
    
    init(from result: ComparisonResult) {
        timestamp = ISO8601DateFormatter().string(from: result.timestamp)
        audio = result.audioResults.map { AudioResultJSON(from: $0) }
        vision = result.visionResults.map { VisionResultJSON(from: $0) }
        endToEnd = result.endToEndResults.map { EndToEndJSON(from: $0) }
    }
}

private struct AudioResultJSON: Codable {
    let modelName: String
    let sampleName: String
    let outputRows: Int
    let outputCols: Int
    let min: Float
    let max: Float
    let mean: Float
    let std: Float
    let nanCount: Int
    let forwardTimeMs: Double
    let memoryPeakMB: Double
    let passed: Bool
    
    init(from r: AudioTestResult) {
        modelName = r.modelName
        sampleName = r.sampleName
        outputRows = r.outputShape.0
        outputCols = r.outputShape.1
        min = r.min
        max = r.max
        mean = r.mean
        std = r.std
        nanCount = r.nanCount
        forwardTimeMs = r.forwardTimeMs
        memoryPeakMB = r.memoryPeakMB
        passed = r.passed
    }
}

private struct VisionResultJSON: Codable {
    let modelName: String
    let sampleName: String
    let outputRows: Int
    let outputCols: Int
    let min: Float
    let max: Float
    let mean: Float
    let std: Float
    let nanCount: Int
    let cosineSimilarity: Float?
    let forwardTimeMs: Double
    let passed: Bool
    
    init(from r: VisionTestResult) {
        modelName = r.modelName
        sampleName = r.sampleName
        outputRows = r.outputShape.0
        outputCols = r.outputShape.1
        min = r.min
        max = r.max
        mean = r.mean
        std = r.std
        nanCount = r.nanCount
        cosineSimilarity = r.cosineSimilarity
        forwardTimeMs = r.forwardTimeMs
        passed = r.passed
    }
}

private struct EndToEndJSON: Codable {
    let modelName: String
    let loadTimeMs: Double
    let audioProcessMs: Double
    let visionProcessMs: Double
    let genSpeedTokPerSec: Double
    let totalTimeMs: Double
    let passed: Bool
    
    init(from r: EndToEndResult) {
        modelName = r.modelName
        loadTimeMs = r.loadTimeMs
        audioProcessMs = r.audioProcessMs
        visionProcessMs = r.visionProcessMs
        genSpeedTokPerSec = r.genSpeedTokPerSec
        totalTimeMs = r.totalTimeMs
        passed = r.passed
    }
}