#import "IMCMCPServer.h"
#import "IMCViewController.h"
#import <arpa/inet.h>
#import <netinet/in.h>
#import <sys/socket.h>
#import <unistd.h>

@interface IMCMCPServer ()
@property (nonatomic, weak) IMCViewController *viewController;
@property (nonatomic, assign) NSInteger port;
@property (nonatomic, assign) BOOL running;
@end

@implementation IMCMCPServer

- (instancetype)initWithViewController:(IMCViewController *)viewController port:(NSInteger)port {
    self = [super init];
    if (self) {
        _viewController = viewController;
        _port = port;
    }
    return self;
}

- (void)start {
    if (self.running) {
        return;
    }
    self.running = YES;
    [NSThread detachNewThreadSelector:@selector(serverLoop) toTarget:self withObject:nil];
}

- (void)serverLoop {
    @autoreleasepool {
        int serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket < 0) {
            return;
        }

        int yes = 1;
        setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

        struct sockaddr_in address;
        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = INADDR_ANY;
        address.sin_port = htons((uint16_t)self.port);

        if (bind(serverSocket, (struct sockaddr *)&address, sizeof(address)) < 0) {
            close(serverSocket);
            return;
        }

        if (listen(serverSocket, 8) < 0) {
            close(serverSocket);
            return;
        }

        while (self.running) {
            @autoreleasepool {
                struct sockaddr_in clientAddress;
                socklen_t clientLength = sizeof(clientAddress);
                int clientSocket = accept(serverSocket,
                                          (struct sockaddr *)&clientAddress,
                                          &clientLength);
                if (clientSocket >= 0) {
                    [self handleClientSocket:clientSocket];
                    close(clientSocket);
                }
            }
        }

        close(serverSocket);
    }
}

- (void)handleClientSocket:(int)clientSocket {
    NSMutableData *requestData = [NSMutableData data];
    char buffer[2048];

    while (requestData.length < 1024 * 1024) {
        ssize_t count = recv(clientSocket, buffer, sizeof(buffer), 0);
        if (count <= 0) {
            break;
        }
        [requestData appendBytes:buffer length:(NSUInteger)count];

        NSString *partial = [[NSString alloc] initWithData:requestData encoding:NSUTF8StringEncoding];
        NSRange separator = [partial rangeOfString:@"\r\n\r\n"];
        if (separator.location == NSNotFound) {
            continue;
        }

        NSInteger contentLength = [self contentLengthFromHeader:partial];
        NSUInteger bodyStart = separator.location + separator.length;
        if (requestData.length >= bodyStart + MAX(contentLength, 0)) {
            break;
        }
    }

    NSString *request = [[NSString alloc] initWithData:requestData encoding:NSUTF8StringEncoding];
    NSDictionary *responseObject = [self responseObjectForHTTPRequest:request];
    NSData *responseJSON = [NSJSONSerialization dataWithJSONObject:responseObject options:0 error:nil];
    if (!responseJSON) {
        responseJSON = [@"{}" dataUsingEncoding:NSUTF8StringEncoding];
    }

    NSString *status = responseObject[@"httpStatus"] ?: @"200 OK";
    NSMutableDictionary *bodyObject = [responseObject mutableCopy];
    [bodyObject removeObjectForKey:@"httpStatus"];
    responseJSON = [NSJSONSerialization dataWithJSONObject:bodyObject options:0 error:nil];

    NSMutableData *responseData = [NSMutableData data];
    NSString *header = [NSString stringWithFormat:
                        @"HTTP/1.1 %@\r\nContent-Type: application/json\r\nContent-Length: %lu\r\nConnection: close\r\n\r\n",
                        status,
                        (unsigned long)responseJSON.length];
    [responseData appendData:[header dataUsingEncoding:NSUTF8StringEncoding]];
    [responseData appendData:responseJSON];
    send(clientSocket, responseData.bytes, responseData.length, 0);
}

- (NSDictionary *)responseObjectForHTTPRequest:(NSString *)request {
    if (![request hasPrefix:@"POST /api/mcp "]) {
        return @{@"httpStatus": @"404 Not Found",
                 @"jsonrpc": @"2.0",
                 @"error": @{@"code": @-32601, @"message": @"POST /api/mcp required"},
                 @"id": [NSNull null]};
    }

    NSRange separator = [request rangeOfString:@"\r\n\r\n"];
    if (separator.location == NSNotFound) {
        return [self parseError:@"Missing HTTP body."];
    }

    NSString *body = [request substringFromIndex:separator.location + separator.length];
    NSData *bodyData = [body dataUsingEncoding:NSUTF8StringEncoding];
    NSError *jsonError = nil;
    NSDictionary *rpcRequest = [NSJSONSerialization JSONObjectWithData:bodyData options:0 error:&jsonError];
    if (jsonError || ![rpcRequest isKindOfClass:NSDictionary.class]) {
        return [self parseError:@"Invalid JSON-RPC object."];
    }

    return [self.viewController handleRPCRequest:rpcRequest];
}

- (NSDictionary *)parseError:(NSString *)message {
    return @{@"jsonrpc": @"2.0",
             @"error": @{@"code": @-32700, @"message": message ?: @"Parse error"},
             @"id": [NSNull null]};
}

- (NSInteger)contentLengthFromHeader:(NSString *)request {
    NSArray *lines = [request componentsSeparatedByString:@"\r\n"];
    for (NSString *line in lines) {
        if ([[line lowercaseString] hasPrefix:@"content-length:"]) {
            return [[line substringFromIndex:15] integerValue];
        }
    }
    return 0;
}

@end