//
//  SanaDiffusionSession.mm
//  MNNLLMiOS
//
//  Created by 游薪渝(揽清) on 2025/01/19.
//  Copyright © 2025 MNN. All rights reserved.
//

#import "SanaDiffusionSession.h"

#import <Foundation/Foundation.h>
#import <UIKit/UIKit.h>

#include <MNN/diffusion/diffusion.hpp>
#include <MNN/diffusion/sana_llm.hpp>
#include <MNN/expr/Module.hpp>
#include <MNN/expr/ExprCreator.hpp>
#include <MNN/expr/Executor.hpp>
#include <MNN/expr/ExecutorScope.hpp>
#include <MNN/Interpreter.hpp>
#include <MNN/AutoTime.hpp>

#include <random>
#include <algorithm>
#include <fstream>

using namespace MNN;
using namespace MNN::Express;
using namespace MNN::Transformer;
using namespace MNN::DIFFUSION;
using namespace CV;

#pragma mark - Private Interface

@implementation SanaDiffusionSession {
    /// SanaLlm instance for processing text prompts to embeddings (from sana_llm.hpp).
    std::unique_ptr<SanaLlm> mSanaLlm;
    
    /// Diffusion model instance using SANA_DIFFUSION type.
    std::shared_ptr<Diffusion> mDiffusion;
    
    /// Path to the model directory.
    NSString *mModelPath;
    
    /// Memory mode for Diffusion model (0 = default).
    int mMemoryMode;
    
    /// Flag indicating whether the model has been loaded.
    BOOL _isModelLoaded;
    
    /// Flag indicating whether a style transfer is in progress.
    BOOL _isProcessing;
}

#pragma mark - Class Methods

+ (NSString *)defaultGhibliPrompt {
    return @"Convert to a Ghibli-style illustration: soft contrast, warm tones, slight linework, keep the scene consistent.";
}

#pragma mark - Initialization

- (instancetype)initWithModelPath:(NSString *)path completion:(SanaCompletionHandler)completion {
    self = [super init];
    if (self) {
        mModelPath = path;
        mMemoryMode = 1;
        _isModelLoaded = NO;
        _isProcessing = NO;
        
        // Load model asynchronously on background thread
        // Threading strategy: All model loading and inference on background thread
        // to keep UI responsive. Accept UIApplication warning for better UX.
        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
            // Step 1: Load LLM on background thread
            BOOL llmSuccess = [self loadLLMModel];
            if (!llmSuccess) {
                dispatch_async(dispatch_get_main_queue(), ^{
                    self->_isModelLoaded = NO;
                    if (completion) {
                        completion(NO);
                    }
                });
                return;
            }
            
            // Step 2: Load Diffusion on background thread
            BOOL diffusionSuccess = [self loadDiffusionModel];
            
            dispatch_async(dispatch_get_main_queue(), ^{
                self->_isModelLoaded = diffusionSuccess;
                if (completion) {
                    completion(diffusionSuccess);
                }
            });
        });
    }
    return self;
}

#pragma mark - Model Loading

/// Loads the SanaLlm model from disk using sana_llm.hpp. Can be called from background thread.
/// @return YES if SanaLlm loaded successfully, NO otherwise.
- (BOOL)loadLLMModel {
    @try {
        NSLog(@"SanaDiffusionSession: Starting SanaLlm loading from %@", mModelPath);
        
        // Load SanaLlm using sana_llm.hpp - it handles LLM and meta_queries loading internally
        NSString *llmPath = [mModelPath stringByAppendingPathComponent:@"llm"];
        NSLog(@"SanaDiffusionSession: Loading SanaLlm from %@", llmPath);
        
        mSanaLlm = std::make_unique<SanaLlm>([llmPath UTF8String]);
        if (!mSanaLlm) {
            NSLog(@"SanaDiffusionSession: Failed to create SanaLlm");
            return NO;
        }
        
        NSLog(@"SanaDiffusionSession: SanaLlm loading complete");
        return YES;
        
    } @catch (NSException *exception) {
        NSLog(@"SanaDiffusionSession: Exception during SanaLlm loading: %@", exception.reason);
        return NO;
    }
}

/// Loads the Diffusion model from disk. Must be called from main thread (Metal requirement).
/// @return YES if Diffusion loaded successfully, NO otherwise.
- (BOOL)loadDiffusionModel {
    @try {
        // Load Diffusion model with SANA_DIFFUSION type
        // Note: Must be on main thread because Metal backend calls [UIApplication applicationState]
        NSLog(@"SanaDiffusionSession: Loading Sana Diffusion model");
        Diffusion* rawDiffusion = Diffusion::createDiffusion(
            [mModelPath UTF8String],
            DiffusionModelType::SANA_DIFFUSION,
            MNNForwardType::MNN_FORWARD_METAL,
            mMemoryMode
        );
        
        if (!rawDiffusion) {
            NSLog(@"SanaDiffusionSession: Failed to create Diffusion");
            return NO;
        }
        
        mDiffusion = std::shared_ptr<Diffusion>(rawDiffusion);
        if (!mDiffusion->load()) {
            NSLog(@"SanaDiffusionSession: Failed to load Diffusion model");
            return NO;
        }
        
        NSLog(@"SanaDiffusionSession: Model loading complete");
        return YES;
        
    } @catch (NSException *exception) {
        NSLog(@"SanaDiffusionSession: Exception during Diffusion loading: %@", exception.reason);
        return NO;
    }
}

#pragma mark - LLM Processing

/// Processes a single text prompt through SanaLlm to generate embeddings for diffusion.
/// Uses sana_llm.hpp's SanaLlm::process() method directly.
///
/// @param prompt The text prompt to process.
/// @return VARP containing the processed embeddings [1, 256, hidden_size], or nullptr on failure.
- (VARP)processSinglePrompt:(NSString *)prompt {
    if (!mSanaLlm) {
        NSLog(@"SanaDiffusionSession: ERROR - mSanaLlm is not initialized!");
        return nullptr;
    }
    
    NSLog(@"SanaDiffusionSession: Processing prompt via SanaLlm: %@", prompt);
    
    // Use SanaLlm::process() directly - it handles all the formatting, tokenization, and LLM forward
    VARP result = mSanaLlm->process([prompt UTF8String]);
    
    if (result.get() == nullptr) {
        NSLog(@"SanaDiffusionSession: ERROR - SanaLlm::process() returned nullptr");
        return nullptr;
    }
    
    // Debug output
    auto resultInfo = result->getInfo();
    if (resultInfo) {
        NSLog(@"SanaDiffusionSession: SanaLlm output shape: [%d, %d, %d]",
              resultInfo->dim[0], resultInfo->dim[1], resultInfo->dim[2]);
    }
    
    return result;
}

#pragma mark - Style Transfer

- (void)runStyleTransferWithInputImage:(NSString *)inputImagePath
                                prompt:(NSString *)prompt
                            outputPath:(NSString *)outputPath
                            iterations:(int)iterations
                                  seed:(int)seed
                      progressCallback:(SanaProgressHandler)progressCallback
                            completion:(SanaStyleTransferCompletion)completion {
    
    // Validate model state
    if (!_isModelLoaded) {
        if (completion) {
            completion(NO, @"Model not loaded", 0);
        }
        return;
    }
    
    if (_isProcessing) {
        if (completion) {
            completion(NO, @"Already processing", 0);
        }
        return;
    }
    
    _isProcessing = YES;
    
    // Process on background thread to keep UI responsive.
    // Note: This will trigger "[UIApplication applicationState] must be used from main thread only"
    // warning from MNN Metal backend. This is expected and acceptable - see threading strategy
    // comment in initWithModelPath:completion: for details.
    dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
        @try {
            // Normalize image orientation if needed
            NSString *normalizedInputPath = [self normalizeImageOrientation:inputImagePath];
            if (!normalizedInputPath) {
                normalizedInputPath = inputImagePath;
            }

            // Record total start time
            NSDate *totalStartTime = [NSDate date];
            NSDate *llmStartTime = [NSDate date];
            
            // Stage 1: Process prompt with LLM
            dispatch_async(dispatch_get_main_queue(), ^{
                if (progressCallback) {
                    progressCallback(5, @"Processing prompt...");
                }
            });
            
            NSLog(@"SanaDiffusionSession: Processing prompt: %@", prompt);
            VARP llmOutput = [self processSinglePrompt:prompt];
            
            NSTimeInterval llmDuration = [[NSDate date] timeIntervalSinceDate:llmStartTime] * 1000;
            NSLog(@"SanaDiffusionSession: LLM processing time: %.2f ms", llmDuration);
            
            if (llmOutput.get() == nullptr) {
                dispatch_async(dispatch_get_main_queue(), ^{
                    self->_isProcessing = NO;
                    if (completion) {
                        completion(NO, @"LLM processing failed", 0);
                    }
                });
                return;
            }
            
            dispatch_async(dispatch_get_main_queue(), ^{
                if (progressCallback) {
                    progressCallback(15, @"Running diffusion...");
                }
            });
            
            // Stage 2: Run diffusion pipeline with LLM embeddings
            NSDate *diffusionStartTime = [NSDate date];
            
            // The pre-built MNN.framework's SanaDiffusion::run() does not invoke the
            // C++ progressCallback during denoising. Work around this by using a
            // dispatch timer on the main queue that fires while run() blocks the
            // background thread, giving the user periodic progress updates.
            __block int timerProgress = 15;
            int progressPerTick = MAX(1, 80 / iterations);  // 80% span / iterations
            
            dispatch_source_t progressTimer = dispatch_source_create(
                DISPATCH_SOURCE_TYPE_TIMER, 0, 0, dispatch_get_main_queue());
            // Fire every 2 seconds (roughly matches per-step duration)
            dispatch_source_set_timer(progressTimer,
                dispatch_time(DISPATCH_TIME_NOW, 2 * NSEC_PER_SEC),
                2 * NSEC_PER_SEC, 0.5 * NSEC_PER_SEC);
            dispatch_source_set_event_handler(progressTimer, ^{
                timerProgress += progressPerTick;
                if (timerProgress > 90) timerProgress = 90;  // cap before completion
                if (progressCallback) {
                    progressCallback(timerProgress, @"Running diffusion...");
                }
            });
            dispatch_resume(progressTimer);
            
            // Also keep the C++ callback in case a future framework rebuild enables it
            auto diffusionProgressCallback = [progressCallback](int step) {
                dispatch_async(dispatch_get_main_queue(), ^{
                    if (progressCallback) {
                        int progress = 15 + (step * 80 / 100);
                        progressCallback(progress, @"Running diffusion...");
                    }
                });
            };
            
            // Debug: Print llmOutput tensor info
            auto llmOutputInfo = llmOutput->getInfo();
            if (llmOutputInfo) {
                NSMutableString *dimStr = [NSMutableString stringWithString:@"["];
                for (size_t i = 0; i < llmOutputInfo->dim.size(); i++) {
                    [dimStr appendFormat:@"%d", llmOutputInfo->dim[i]];
                    if (i < llmOutputInfo->dim.size() - 1) {
                        [dimStr appendString:@", "];
                    }
                }
                [dimStr appendString:@"]"];
                NSLog(@"SanaDiffusionSession: llmOutput shape: %@, size: %zu, order: %d, type: %d",
                      dimStr,
                      llmOutputInfo->size,
                      (int)llmOutputInfo->order,
                      (int)llmOutputInfo->type.code);
            } else {
                NSLog(@"SanaDiffusionSession: llmOutput info is NULL!");
            }
            
            // Use new unified diffusion interface
            bool success = self->mDiffusion->run(
                llmOutput,
                "img2img",                       // mode: image editing
                [normalizedInputPath UTF8String], // input image path
                [outputPath UTF8String],         // output image path
                512,                             // width
                512,                             // height
                iterations,                      // iterNum
                seed,                            // randomSeed
                false,                           // use_cfg
                4.5f,                            // cfg_scale
                diffusionProgressCallback
            );
            
            // Cancel the progress timer now that run() has returned
            dispatch_source_cancel(progressTimer);
            
            
            NSTimeInterval diffusionDuration = [[NSDate date] timeIntervalSinceDate:diffusionStartTime] * 1000;
            NSTimeInterval totalDuration = [[NSDate date] timeIntervalSinceDate:totalStartTime] * 1000;
            
            NSLog(@"SanaDiffusionSession: Diffusion time: %.2f ms", diffusionDuration);
            NSLog(@"SanaDiffusionSession: Total time: %.2f ms", totalDuration);
            
            // Clean up temporary normalized image if it was created
            if (![normalizedInputPath isEqualToString:inputImagePath]) {
                [[NSFileManager defaultManager] removeItemAtPath:normalizedInputPath error:nil];
            }

            // Save benchmark results to file
            [self saveBenchmarkResult:@{
                @"timestamp": [NSDate date],
                @"iterations": @(iterations),
                @"seed": @(seed),
                @"llm_time_ms": @(llmDuration),
                @"diffusion_time_ms": @(diffusionDuration),
                @"total_time_ms": @(totalDuration),
                @"success": @(success),
                @"input_image": inputImagePath,
                @"output_image": outputPath
            }];
            
            // Report completion on main thread
            dispatch_async(dispatch_get_main_queue(), ^{
                self->_isProcessing = NO;
                
                // Note: Do NOT call progressCallback(100) here.
                // The completion handler sets the final message directly,
                // so calling progressCallback would cause a brief flash of
                // an inconsistent message before completion replaces it.
                
                if (completion) {
                    if (success) {
                        completion(YES, nil, totalDuration);
                    } else {
                        completion(NO, @"Diffusion processing failed", totalDuration);
                    }
                }
            });
            
        } @catch (NSException *exception) {
            dispatch_async(dispatch_get_main_queue(), ^{
                self->_isProcessing = NO;
                if (completion) {
                    completion(NO, [NSString stringWithFormat:@"Exception: %@", exception.reason], 0);
                }
            });
        }
    });
}

#pragma mark - Image Orientation Normalization

/// Normalizes the image orientation based on EXIF data.
/// @param imagePath Path to the input image.
/// @return Path to the normalized image, or nil if failed/not needed.
- (NSString *)normalizeImageOrientation:(NSString *)imagePath {
    UIImage *image = [UIImage imageWithContentsOfFile:imagePath];
    if (!image) {
        NSLog(@"SanaDiffusionSession: Failed to load image for normalization: %@", imagePath);
        return nil;
    }
    
    // Check if orientation is already Up
    if (image.imageOrientation == UIImageOrientationUp) {
        return imagePath;
    }
    
    NSLog(@"SanaDiffusionSession: Normalizing image orientation from %ld", (long)image.imageOrientation);
    
    // Drawing into a context redraws the image with "Up" orientation
    UIGraphicsBeginImageContextWithOptions(image.size, NO, image.scale);
    [image drawInRect:CGRectMake(0, 0, image.size.width, image.size.height)];
    UIImage *normalizedImage = UIGraphicsGetImageFromCurrentImageContext();
    UIGraphicsEndImageContext();
    
    if (!normalizedImage) {
        NSLog(@"SanaDiffusionSession: Failed to normalize image");
        return nil;
    }
    
    // Save normalized image to a temporary path
    NSString *tempPath = [NSTemporaryDirectory() stringByAppendingPathComponent:[NSString stringWithFormat:@"normalized_%@", [imagePath lastPathComponent]]];
    NSData *data = UIImageJPEGRepresentation(normalizedImage, 0.9);
    if ([data writeToFile:tempPath atomically:YES]) {
        NSLog(@"SanaDiffusionSession: Normalized image saved to %@", tempPath);
        return tempPath;
    } else {
        NSLog(@"SanaDiffusionSession: Failed to save normalized image");
        return nil;
    }
}

#pragma mark - Benchmark

/// Saves benchmark result to a JSON file in the Documents directory.
/// @param result Dictionary containing benchmark metrics.
- (void)saveBenchmarkResult:(NSDictionary *)result {
    NSFileManager *fileManager = [NSFileManager defaultManager];
    NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES);
    NSString *documentsDirectory = [paths firstObject];
    NSString *benchmarkFile = [documentsDirectory stringByAppendingPathComponent:@"sana_diffusion_benchmark.json"];
    
    // Read existing results or create new array
    NSMutableArray *results = [NSMutableArray array];
    if ([fileManager fileExistsAtPath:benchmarkFile]) {
        NSData *existingData = [NSData dataWithContentsOfFile:benchmarkFile];
        if (existingData) {
            NSError *error = nil;
            NSArray *existingResults = [NSJSONSerialization JSONObjectWithData:existingData options:0 error:&error];
            if (existingResults && !error) {
                [results addObjectsFromArray:existingResults];
            }
        }
    }
    
    // Format timestamp as string
    NSDateFormatter *formatter = [[NSDateFormatter alloc] init];
    [formatter setDateFormat:@"yyyy-MM-dd HH:mm:ss"];
    NSMutableDictionary *formattedResult = [result mutableCopy];
    if (result[@"timestamp"]) {
        formattedResult[@"timestamp"] = [formatter stringFromDate:result[@"timestamp"]];
    }
    
    // Add new result
    [results addObject:formattedResult];
    
    // Write back to file
    NSError *writeError = nil;
    NSData *jsonData = [NSJSONSerialization dataWithJSONObject:results options:NSJSONWritingPrettyPrinted error:&writeError];
    if (jsonData && !writeError) {
        [jsonData writeToFile:benchmarkFile atomically:YES];
        NSLog(@"SanaDiffusionSession: Benchmark saved to %@", benchmarkFile);
    } else {
        NSLog(@"SanaDiffusionSession: Failed to save benchmark: %@", writeError);
    }
}

/// Returns the path to the benchmark results file.
+ (NSString *)benchmarkFilePath {
    NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES);
    NSString *documentsDirectory = [paths firstObject];
    return [documentsDirectory stringByAppendingPathComponent:@"sana_diffusion_benchmark.json"];
}

#pragma mark - Properties

- (BOOL)isModelLoaded {
    return _isModelLoaded;
}

- (BOOL)isProcessing {
    return _isProcessing;
}

#pragma mark - Memory Management

- (void)dealloc {
    mSanaLlm.reset();
    mDiffusion.reset();
    NSLog(@"SanaDiffusionSession deallocated");
}

@end