import Foundation
import NIOCore
import NIOPosix
import NIOHTTP1

/// Handles incoming proxy requests:
/// - HTTP CONNECT -> TCP tunnel (GlueHandler passthrough or MITM)
/// - Plain HTTP -> forward with capture
final class ConnectHandler: ChannelInboundHandler, RemovableChannelHandler {
    typealias InboundIn = HTTPServerRequestPart
    typealias OutboundOut = HTTPServerResponsePart

    private let trafficRepo: TrafficRepository
    private let runtimeStatusRepo = RuntimeStatusRepository()

    private var pendingConnectHead: HTTPRequestHead?
    private var pendingConnectBytes: [ByteBuffer] = []

    private var pendingHead: HTTPRequestHead?
    private var pendingBody: [ByteBuffer] = []
    private var pendingEnd: HTTPHeaders?

    init(trafficRepo: TrafficRepository) {
        self.trafficRepo = trafficRepo
    }

    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let part = unwrapInboundIn(data)

        switch part {
        case .head(let head):
            if head.method == .CONNECT {
                ProxyLogger.connect.info("CONNECT \(head.uri)")
                pendingConnectHead = head
                pendingConnectBytes.removeAll()
            } else {
                ProxyLogger.connect.info("HTTP \(head.method.rawValue) \(head.uri)")
                pendingHead = head
                pendingBody.removeAll()
                pendingEnd = nil
            }

        case .body(let buffer):
            if pendingConnectHead != nil {
                pendingConnectBytes.append(buffer)
            } else {
                pendingBody.append(buffer)
            }

        case .end(let trailers):
            if let connectHead = pendingConnectHead {
                let bufferedBytes = pendingConnectBytes
                pendingConnectHead = nil
                pendingConnectBytes.removeAll()
                handleConnect(context: context, head: connectHead, initialBuffers: bufferedBytes)
                return
            }

            if pendingHead != nil {
                pendingEnd = trailers
                handleHTTPRequest(context: context)
            }
        }
    }

    // MARK: - CONNECT

    private func handleConnect(
        context: ChannelHandlerContext,
        head: HTTPRequestHead,
        initialBuffers: [ByteBuffer]
    ) {
        let components = head.uri.split(separator: ":")
        let originalHost = String(components[0])
        let port = components.count > 1 ? Int(components[1]) ?? 443 : 443
        let connectURL = "https://\(originalHost):\(port)"

        if let blockAction = RulesEngine.checkBlockList(url: connectURL, method: "CONNECT"),
           blockAction != .hideOnly {
            ProxyLogger.connect.info("BLOCKED \(originalHost) action=\(blockAction.rawValue)")
            if blockAction == .blockAndDisplay {
                var traffic = CapturedTraffic(
                    domain: originalHost, url: connectURL, method: "CONNECT", scheme: "https",
                    statusCode: 403, statusText: "Blocked",
                    startedAt: Date().timeIntervalSince1970,
                    completedAt: Date().timeIntervalSince1970, durationMs: 0, isSslDecrypted: false
                )
                do {
                    try trafficRepo.insert(&traffic)
                    IPCManager.shared.post(.newTrafficCaptured)
                } catch {
                    ProxyLogger.db.error("DB insert blocked traffic failed: \(error.localizedDescription)")
                }
            }

            let responseHead = HTTPResponseHead(version: .http1_1, status: .forbidden)
            context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
            context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil)
            context.close(promise: nil)
            return
        }

        let upstreamHost = RulesEngine.checkDNSSpoof(domain: originalHost) ?? originalHost
        let shouldMITM = shouldInterceptSSL(domain: originalHost)
        let shouldHide = shouldHideConnect(url: connectURL, host: originalHost)
        ProxyLogger.connect.info("=== CONNECT original=\(originalHost) upstream=\(upstreamHost):\(port) mitm=\(shouldMITM) ===")

        if shouldMITM {
            upgradeToMITM(
                context: context,
                originalHost: originalHost,
                upstreamHost: upstreamHost,
                port: port,
                initialBuffers: initialBuffers
            )
            return
        }

        ClientBootstrap(group: context.eventLoop)
            .channelOption(.socketOption(.so_reuseaddr), value: 1)
            .channelOption(.autoRead, value: false)
            .connect(host: upstreamHost, port: port)
            .whenComplete { result in
                switch result {
                case .success(let remoteChannel):
                    ProxyLogger.connect.info("Upstream connected to \(upstreamHost):\(port), upgrading to passthrough")
                    self.upgradeToPassthrough(
                        context: context,
                        remoteChannel: remoteChannel,
                        originalHost: originalHost,
                        upstreamHost: upstreamHost,
                        port: port,
                        initialBuffers: initialBuffers,
                        isHidden: shouldHide
                    )

                case .failure(let error):
                    ProxyLogger.connect.error("Upstream connect FAILED \(upstreamHost):\(port): \(error.localizedDescription)")
                    self.runtimeStatusRepo.update {
                        $0.lastConnectError = "CONNECT \(originalHost): \(error.localizedDescription)"
                    }
                    let responseHead = HTTPResponseHead(version: .http1_1, status: .badGateway)
                    context.write(self.wrapOutboundOut(.head(responseHead)), promise: nil)
                    context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
                    context.close(promise: nil)
                }
            }
    }

    private func shouldInterceptSSL(domain: String) -> Bool {
        let sslEnabled = IPCManager.shared.isSSLProxyingEnabled
        let hasCA = CertificateManager.shared.hasCA
        ProxyLogger.connect.info("shouldInterceptSSL(\(domain)): sslEnabled=\(sslEnabled) hasCA=\(hasCA)")

        // Write diagnostic info so the app can display what the extension sees
        runtimeStatusRepo.update {
            $0.caFingerprint = CertificateManager.shared.caFingerprint
            $0.lastConnectError = "SSL check: domain=\(domain) sslEnabled=\(sslEnabled) hasCA=\(hasCA)"
        }

        guard sslEnabled else {
            ProxyLogger.connect.info("SSL proxying DISABLED globally — skipping MITM")
            runtimeStatusRepo.update {
                $0.lastMITMError = "SSL proxying disabled (sslEnabled=false in DB)"
            }
            return false
        }
        guard hasCA else {
            ProxyLogger.connect.info("Shared CA unavailable in extension — skipping MITM")
            runtimeStatusRepo.update {
                $0.lastMITMError = "No CA in extension (hasCA=false)"
            }
            return false
        }

        // Check if domain was auto-detected as using SSL pinning
        if PinnedDomainRepository().isPinned(domain: domain) {
            ProxyLogger.connect.info("SSL PINNED (auto-detected): \(domain) — using passthrough")
            runtimeStatusRepo.update {
                $0.lastMITMError = "Pinned domain (auto-fallback): \(domain)"
            }
            return false
        }

        let rulesRepo = RulesRepository()
        do {
            let entries = try rulesRepo.fetchAllSSLEntries()
            let includeCount = entries.filter(\.isInclude).count
            let excludeCount = entries.filter { !$0.isInclude }.count
            let patterns = entries.map { "\($0.isInclude ? "+" : "-")\($0.domainPattern)" }.joined(separator: ", ")
            ProxyLogger.connect.info("SSL entries: \(entries.count) (include=\(includeCount) exclude=\(excludeCount)) patterns=[\(patterns)]")

            runtimeStatusRepo.update {
                $0.lastConnectError = "SSL rules: \(entries.count) entries [\(patterns)] checking domain=\(domain)"
            }

            for entry in entries where !entry.isInclude {
                if WildcardMatcher.matches(domain, pattern: entry.domainPattern) {
                    ProxyLogger.connect.debug("SSL EXCLUDED by pattern: \(entry.domainPattern)")
                    return false
                }
            }
            for entry in entries where entry.isInclude {
                if WildcardMatcher.matches(domain, pattern: entry.domainPattern) {
                    ProxyLogger.connect.info("SSL INCLUDED by pattern: \(entry.domainPattern) -> MITM ON")
                    runtimeStatusRepo.update {
                        $0.lastMITMError = nil
                        $0.lastConnectError = "MITM enabled for \(domain) via pattern \(entry.domainPattern)"
                    }
                    return true
                }
            }
        } catch {
            ProxyLogger.connect.error("SSL list fetch failed: \(error.localizedDescription)")
            runtimeStatusRepo.update {
                $0.lastMITMError = "SSL list DB error: \(error.localizedDescription)"
            }
        }

        ProxyLogger.connect.debug("SSL: no matching rule for \(domain)")
        return false
    }

    private func shouldHideConnect(url: String, host: String) -> Bool {
        if let blockAction = RulesEngine.checkBlockList(url: url, method: "CONNECT"), blockAction == .hideOnly {
            return true
        }
        return IPCManager.shared.hideSystemTraffic && SystemTrafficFilter.isSystemDomain(host)
    }

    private func upgradeToMITM(
        context: ChannelHandlerContext,
        originalHost: String,
        upstreamHost: String,
        port: Int,
        initialBuffers: [ByteBuffer]
    ) {
        let channel = context.channel

        channel.setOption(.autoRead, value: false).flatMap {
            self.upgradeClientChannelToRaw(channel)
        }.flatMap {
            channel.pipeline.addHandler(
                MITMHandler(
                    originalHost: originalHost,
                    upstreamHost: upstreamHost,
                    port: port,
                    trafficRepo: self.trafficRepo
                )
            )
        }.flatMap {
            self.sendConnectEstablished(on: channel)
        }.whenComplete { result in
            switch result {
            case .success:
                ProxyLogger.connect.info("MITM pipeline ready for \(originalHost):\(port)")
                self.runtimeStatusRepo.update {
                    $0.lastConnectError = nil
                    $0.lastMITMError = nil
                }
                channel.setOption(.autoRead, value: true).whenComplete { _ in
                    channel.read()
                    for buffer in initialBuffers {
                        channel.pipeline.fireChannelRead(NIOAny(buffer))
                    }
                }

            case .failure(let error):
                ProxyLogger.connect.error("MITM upgrade FAILED for \(originalHost):\(port): \(error.localizedDescription)")
                self.runtimeStatusRepo.update {
                    $0.lastMITMError = "MITM setup \(originalHost): \(error.localizedDescription)"
                }
                channel.close(promise: nil)
            }
        }
    }

    private func upgradeToPassthrough(
        context: ChannelHandlerContext,
        remoteChannel: Channel,
        originalHost: String,
        upstreamHost: String,
        port: Int,
        initialBuffers: [ByteBuffer],
        isHidden: Bool
    ) {
        let channel = context.channel
        let localGlue = GlueHandler()
        let remoteGlue = GlueHandler()
        localGlue.partner = remoteGlue
        remoteGlue.partner = localGlue

        channel.setOption(.autoRead, value: false).flatMap {
            self.upgradeClientChannelToRaw(channel)
        }.flatMap {
            remoteChannel.pipeline.addHandler(remoteGlue)
        }.flatMap {
            channel.pipeline.addHandler(localGlue)
        }.flatMap {
            self.sendConnectEstablished(on: channel)
        }.whenComplete { result in
            switch result {
            case .success:
                ProxyLogger.connect.info("Passthrough tunnel ready for \(originalHost):\(port) via \(upstreamHost)")
                self.runtimeStatusRepo.update {
                    $0.lastConnectError = nil
                }
                self.recordConnectTraffic(host: originalHost, port: port, isHidden: isHidden)

                for buffer in initialBuffers {
                    remoteChannel.write(NIOAny(buffer), promise: nil)
                }
                remoteChannel.flush()

                channel.setOption(.autoRead, value: true).whenComplete { _ in
                    channel.read()
                }
                remoteChannel.setOption(.autoRead, value: true).whenComplete { _ in
                    remoteChannel.read()
                }

            case .failure(let error):
                ProxyLogger.connect.error("Passthrough upgrade FAILED for \(originalHost):\(port): \(error.localizedDescription)")
                self.runtimeStatusRepo.update {
                    $0.lastConnectError = "Passthrough \(originalHost): \(error.localizedDescription)"
                }
                channel.close(promise: nil)
                remoteChannel.close(promise: nil)
            }
        }
    }

    private func upgradeClientChannelToRaw(_ channel: Channel) -> EventLoopFuture<Void> {
        removeHandler(ByteToMessageHandler<HTTPRequestDecoder>.self, from: channel).flatMap { _ in
            self.removeHandler(HTTPResponseEncoder.self, from: channel)
        }.flatMap { _ in
            channel.pipeline.removeHandler(self)
        }
    }

    private func sendConnectEstablished(on channel: Channel) -> EventLoopFuture<Void> {
        var buffer = channel.allocator.buffer(capacity: 64)
        buffer.writeString("HTTP/1.1 200 Connection Established\r\n\r\n")
        return channel.writeAndFlush(NIOAny(buffer))
    }

    private func removeHandler<H: RemovableChannelHandler>(_ type: H.Type, from channel: Channel) -> EventLoopFuture<Void> {
        channel.pipeline.handler(type: type).flatMap { handler in
            channel.pipeline.removeHandler(handler)
        }.recover { _ in () }
    }

    // MARK: - Plain HTTP

    private func handleHTTPRequest(context: ChannelHandlerContext) {
        guard let head = pendingHead else { return }

        guard let (host, port, path) = parseHTTPTarget(head: head) else {
            ProxyLogger.connect.error("HTTP: failed to parse target from \(head.uri)")
            let responseHead = HTTPResponseHead(version: .http1_1, status: .badRequest)
            context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
            context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil)
            pendingHead = nil
            pendingBody.removeAll()
            pendingEnd = nil
            return
        }

        let fullURL = "http://\(host)\(path)"
        let method = head.method.rawValue
        let upstreamHost = RulesEngine.checkDNSSpoof(domain: host) ?? host
        ProxyLogger.connect.info("HTTP FORWARD \(method) \(fullURL)")

        if let blockAction = RulesEngine.checkBlockList(url: fullURL, method: method),
           blockAction != .hideOnly {
            ProxyLogger.connect.info("HTTP BLOCKED \(fullURL) action=\(blockAction.rawValue)")
            if blockAction == .blockAndDisplay {
                var traffic = CapturedTraffic(
                    domain: host, url: fullURL, method: method, scheme: "http",
                    statusCode: 403, statusText: "Blocked",
                    startedAt: Date().timeIntervalSince1970,
                    completedAt: Date().timeIntervalSince1970, durationMs: 0, isSslDecrypted: false
                )
                do {
                    try trafficRepo.insert(&traffic)
                    IPCManager.shared.post(.newTrafficCaptured)
                } catch {
                    ProxyLogger.db.error("DB insert failed: \(error.localizedDescription)")
                }
            }
            let responseHead = HTTPResponseHead(version: .http1_1, status: .forbidden)
            context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
            context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil)
            pendingHead = nil
            pendingBody.removeAll()
            pendingEnd = nil
            return
        }

        if let mapRule = RulesEngine.checkMapLocal(url: fullURL, method: method) {
            ProxyLogger.connect.info("MAP LOCAL match for \(fullURL) -> status \(mapRule.responseStatus)")
            let status = HTTPResponseStatus(statusCode: mapRule.responseStatus)
            var headers = decodeHeaders(mapRule.responseHeaders)
            if let ct = mapRule.responseContentType, !ct.isEmpty {
                headers.replaceOrAdd(name: "Content-Type", value: ct)
            }
            let bodyData = mapRule.responseBody
            if let bodyData, !bodyData.isEmpty {
                headers.replaceOrAdd(name: "Content-Length", value: "\(bodyData.count)")
            }
            let responseHead = HTTPResponseHead(version: .http1_1, status: status, headers: headers)
            context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
            if let bodyData, !bodyData.isEmpty {
                var buffer = context.channel.allocator.buffer(capacity: bodyData.count)
                buffer.writeBytes(bodyData)
                context.write(wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil)
            }
            context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil)
            pendingHead = nil
            pendingBody.removeAll()
            pendingEnd = nil
            return
        }

        var upstreamHead = head
        upstreamHead.uri = path
        if !upstreamHead.headers.contains(name: "Host") {
            upstreamHead.headers.add(name: "Host", value: host)
        }

        let captureHandler = HTTPCaptureHandler(trafficRepo: trafficRepo, domain: host, scheme: "http")

        ClientBootstrap(group: context.eventLoop)
            .channelOption(.socketOption(.so_reuseaddr), value: 1)
            .channelInitializer { channel in
                channel.pipeline.addHandler(HTTPRequestEncoder()).flatMap {
                    channel.pipeline.addHandler(ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes)))
                }.flatMap {
                    channel.pipeline.addHandler(captureHandler)
                }.flatMap {
                    channel.pipeline.addHandler(
                        HTTPRelayHandler(clientContext: context, wrapResponse: self.wrapOutboundOut)
                    )
                }
            }
            .connect(host: upstreamHost, port: port)
            .whenComplete { result in
                switch result {
                case .success(let remoteChannel):
                    ProxyLogger.connect.info("HTTP upstream connected to \(upstreamHost):\(port), forwarding request")
                    remoteChannel.write(NIOAny(HTTPClientRequestPart.head(upstreamHead)), promise: nil)
                    for bodyBuffer in self.pendingBody {
                        remoteChannel.write(NIOAny(HTTPClientRequestPart.body(.byteBuffer(bodyBuffer))), promise: nil)
                    }
                    remoteChannel.writeAndFlush(NIOAny(HTTPClientRequestPart.end(self.pendingEnd)), promise: nil)
                    self.pendingHead = nil
                    self.pendingBody.removeAll()
                    self.pendingEnd = nil

                case .failure(let error):
                    ProxyLogger.connect.error("HTTP upstream connect FAILED \(host):\(port): \(error.localizedDescription)")
                    self.runtimeStatusRepo.update {
                        $0.lastConnectError = "HTTP \(fullURL): \(error.localizedDescription)"
                    }
                    let responseHead = HTTPResponseHead(version: .http1_1, status: .badGateway)
                    context.write(self.wrapOutboundOut(.head(responseHead)), promise: nil)
                    context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
                    self.pendingHead = nil
                    self.pendingBody.removeAll()
                    self.pendingEnd = nil
                }
            }
    }

    private func decodeHeaders(_ json: String?) -> HTTPHeaders {
        guard let json,
              let data = json.data(using: .utf8),
              let dict = try? JSONDecoder().decode([String: String].self, from: data) else {
            return HTTPHeaders()
        }

        var headers = HTTPHeaders()
        for (name, value) in dict {
            headers.add(name: name, value: value)
        }
        return headers
    }

    // MARK: - URL Parsing

    private func parseHTTPTarget(head: HTTPRequestHead) -> (host: String, port: Int, path: String)? {
        if head.uri.hasPrefix("http://") || head.uri.hasPrefix("https://") {
            guard let url = URLComponents(string: head.uri) else { return nil }
            let host = url.host ?? ""
            let port = url.port ?? (head.uri.hasPrefix("https") ? 443 : 80)
            var path = url.path.isEmpty ? "/" : url.path
            if let query = url.query { path += "?\(query)" }
            return (host, port, path)
        }
        if let hostHeader = head.headers.first(name: "Host") {
            let parts = hostHeader.split(separator: ":")
            let host = String(parts[0])
            let port = parts.count > 1 ? Int(parts[1]) ?? 80 : 80
            return (host, port, head.uri)
        }
        return nil
    }

    // MARK: - CONNECT traffic recording

    private func recordConnectTraffic(host: String, port: Int, isHidden: Bool) {
        var traffic = CapturedTraffic(
            domain: host, url: "https://\(host):\(port)", method: "CONNECT", scheme: "https",
            statusCode: 200, statusText: "Connection Established",
            startedAt: Date().timeIntervalSince1970, completedAt: Date().timeIntervalSince1970,
            durationMs: 0, isSslDecrypted: false, isHidden: isHidden
        )
        do {
            try trafficRepo.insert(&traffic)
            ProxyLogger.db.debug("Recorded CONNECT \(host) (hidden=\(isHidden))")
        } catch {
            ProxyLogger.db.error("Failed to record CONNECT \(host): \(error.localizedDescription)")
        }
        NotificationThrottle.shared.throttle()
    }
}

// MARK: - HTTPRelayHandler

final class HTTPRelayHandler: ChannelInboundHandler, RemovableChannelHandler {
    typealias InboundIn = HTTPClientResponsePart

    private let clientContext: ChannelHandlerContext
    private let wrapResponse: (HTTPServerResponsePart) -> NIOAny

    init(clientContext: ChannelHandlerContext, wrapResponse: @escaping (HTTPServerResponsePart) -> NIOAny) {
        self.clientContext = clientContext
        self.wrapResponse = wrapResponse
    }

    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let part = unwrapInboundIn(data)
        switch part {
        case .head(let head):
            ProxyLogger.connect.debug("HTTPRelay response: \(head.status.code)")
            clientContext.write(wrapResponse(.head(HTTPResponseHead(version: head.version, status: head.status, headers: head.headers))), promise: nil)
        case .body(let buffer):
            clientContext.write(wrapResponse(.body(.byteBuffer(buffer))), promise: nil)
        case .end(let trailers):
            clientContext.writeAndFlush(wrapResponse(.end(trailers)), promise: nil)
        }
    }

    func channelInactive(context: ChannelHandlerContext) {
        ProxyLogger.connect.debug("HTTPRelay: remote channel inactive")
        clientContext.close(promise: nil)
    }

    func errorCaught(context: ChannelHandlerContext, error: Error) {
        ProxyLogger.connect.error("HTTPRelay error: \(error.localizedDescription)")
        context.close(promise: nil)
        clientContext.close(promise: nil)
    }
}