import Foundation
import Network
import Observation
import SwiftUI
import WebKit

enum Page: Hashable {
    case file(taskId: String, path: String)
    case preview(taskId: String, port: Int)
    case web(taskId: String, id: Int)

    var taskId: String {
        switch self {
        case .file(let taskId, _), .preview(let taskId, _), .web(let taskId, _): taskId
        }
    }

    var isFile: Bool {
        if case .file = self {
            return true
        }
        return false
    }

    @MainActor
    func title(_ store: Store) -> String {
        switch self {
        case .file(_, let path): (path as NSString).lastPathComponent
        case .preview(let taskId, let port): store.forwards.name(taskId, port) ?? "port \(port)"
        case .web: store.pages.title(self)
        }
    }

    @MainActor
    func url(_ store: Store) -> URL? {
        switch self {
        case .file(let taskId, let path): store.pages.url(taskId, path)
        case .preview(let taskId, let port): store.forwards.url(taskId, port)
        case .web: store.pages.existing(self)?.url
        }
    }
}

private let pageExts: Set<String> = ["html", "htm"]

func isPage(_ path: String) -> Bool {
    pageExts.contains((path as NSString).pathExtension.lowercased())
}

struct ShownPages {
    var pages: [Page]
    var current: Page
}

extension Store {
    func openPage(_ page: Page) {
        var shown = shownPages[page.taskId] ?? ShownPages(pages: [], current: page)
        if !shown.pages.contains(page) {
            shown.pages.append(page)
        }
        shown.current = page
        shownPages[page.taskId] = shown
        sheet = Sheet(owner: .task(page.taskId), view: .tab(.page))
    }

    func openWeb(_ taskId: String, _ url: URL) -> Int {
        route("#/task/\(taskId)")
        let page = pages.web(taskId, url)
        openPage(page)
        return pages.tab(page)
    }

    func closePage(_ page: Page) {
        guard var shown = shownPages[page.taskId], let i = shown.pages.firstIndex(of: page) else {
            return
        }
        pages.drop(page)
        shown.pages.remove(at: i)
        if shown.pages.isEmpty {
            shownPages[page.taskId] = nil
            if sheet?.tab == .page {
                sheet = nil
            }
            return
        }
        if shown.current == page {
            shown.current = shown.pages[min(i, shown.pages.count - 1)]
        }
        shownPages[page.taskId] = shown
    }

    func currentPage() -> Page? {
        selectedTaskId.flatMap { shownPages[$0]?.current }
    }

    func pageView(_ page: Page) -> WKWebView? {
        if case .web = page {
            return pages.existing(page)
        }
        return page.url(self).map { pages.view(page, $0) }
    }

    func pageTabs(_ taskId: String) -> String {
        guard let shown = shownPages[taskId] else {
            return "This task's page view has no tabs; open one with show_page."
        }
        return shown.pages.map { page in
            var line = "tab \(pages.tab(page)): \(page.title(self))"
            if let url = pages.existing(page)?.url ?? page.url(self) {
                line += " — \(url.absoluteString)"
            }
            if page == shown.current {
                line += " (showing)"
            }
            return line
        }.joined(separator: "\n")
    }

    func drivenPage(_ taskId: String, _ tab: Int?) throws -> WKWebView {
        guard let shown = shownPages[taskId] else {
            throw ApiError(status: 0, message: "this task's page view has no tabs; open one with show_page")
        }
        var page = shown.current
        if let tab {
            guard let p = shown.pages.first(where: { pages.tab($0) == tab }) else {
                throw ApiError(status: 0, message: "this task's page view has no tab \(tab); drive_page tabs lists them")
            }
            page = p
        }
        guard let view = pageView(page) else {
            throw ApiError(status: 0, message: "tab \(pages.tab(page)) isn't ready yet")
        }
        return view
    }
}

@MainActor
@Observable
final class Pages {
    private(set) var port: Int?
    private(set) var titles: [Page: String] = [:]
    @ObservationIgnored var served: [String: [String: Data]] = [:]
    @ObservationIgnored var popup: ((Page) -> Void)?
    @ObservationIgnored private var listener: PageListener?
    @ObservationIgnored private var views: [Page: WKWebView] = [:]
    @ObservationIgnored private var sources: [Page: URL] = [:]
    @ObservationIgnored private var titleWatches: [Page: NSKeyValueObservation] = [:]
    @ObservationIgnored private var tabs: [Page: Int] = [:]
    @ObservationIgnored private var lastTab = 0
    @ObservationIgnored private let delegate = PageDelegate()

    init(api: API) {
        delegate.popup = { [weak self] opener, config in
            self?.adopt(opener, config)
        }
        listener = PageListener(api: api, ready: { [weak self] port in
            Task { @MainActor in
                self?.port = port
            }
        }, served: { [weak self] taskId, path, body in
            Task { @MainActor in
                self?.served[taskId, default: [:]][path] = body
            }
        })
    }

    func url(_ taskId: String, _ path: String) -> URL? {
        guard let port else {
            return nil
        }
        var c = URLComponents()
        c.scheme = "http"
        c.host = "127.0.0.1"
        c.port = port
        c.path = "/\(taskId)/\(path)"
        return c.url
    }

    func existing(_ page: Page) -> WKWebView? {
        views[page]
    }

    func view(_ page: Page, _ url: URL) -> WKWebView {
        let view = views[page] ?? make(page, WKWebViewConfiguration())
        if sources[page] != url {
            sources[page] = url
            view.load(URLRequest(url: url))
        }
        return view
    }

    func web(_ taskId: String, _ url: URL) -> Page {
        lastTab += 1
        let page = Page.web(taskId: taskId, id: lastTab)
        make(page, WKWebViewConfiguration()).load(URLRequest(url: url))
        return page
    }

    private func adopt(_ opener: WKWebView, _ config: WKWebViewConfiguration) -> WKWebView? {
        guard let from = views.first(where: { $0.value === opener })?.key else {
            return nil
        }
        lastTab += 1
        let page = Page.web(taskId: from.taskId, id: lastTab)
        let view = make(page, config)
        popup?(page)
        return view
    }

    private func make(_ page: Page, _ config: WKWebViewConfiguration) -> WKWebView {
        config.websiteDataStore = .default()
        let view = WKWebView(frame: .zero, configuration: config)
        view.uiDelegate = delegate
        view.navigationDelegate = delegate
        views[page] = view
        if case .web(_, let id) = page {
            tabs[page] = id
            titleWatches[page] = view.observe(\.title) { [weak self] view, _ in
                Task { @MainActor in
                    self?.titles[page] = view.title
                }
            }
        }
        return view
    }

    func drop(_ page: Page) {
        views[page] = nil
        sources[page] = nil
        titleWatches[page] = nil
        titles[page] = nil
        tabs[page] = nil
    }

    func tab(_ page: Page) -> Int {
        if let tab = tabs[page] {
            return tab
        }
        lastTab += 1
        tabs[page] = lastTab
        return lastTab
    }

    func title(_ page: Page) -> String {
        if let title = titles[page], !title.isEmpty {
            return title
        }
        return views[page]?.url?.host() ?? "page"
    }

    func navigate(_ view: WKWebView, _ url: URL) async throws -> String {
        try await delegate.wait(view.load(URLRequest(url: url)))
        return "Loaded \(view.url?.absoluteString ?? url.absoluteString): \(view.title ?? "")"
    }

    func js(_ view: WKWebView, _ code: String) async throws -> String {
        let value = try await view.callAsyncJavaScript(code, contentWorld: .page)
        guard let value, !(value is NSNull) else {
            return value == nil ? "undefined" : "null"
        }
        if let s = value as? String {
            return s
        }
        if let data = try? JSONSerialization.data(withJSONObject: value, options: [.fragmentsAllowed, .sortedKeys]) {
            return String(decoding: data, as: UTF8.self)
        }
        return String(describing: value)
    }

    #if os(macOS)
    func capture(_ view: WKWebView) async throws -> Data {
        guard unsafe view.window != nil, !view.isHiddenOrHasHiddenAncestor else {
            throw ApiError(status: 0, message: "the tab isn't on screen: the human has the page view closed or another tab showing")
        }
        let image = try await view.takeSnapshot(configuration: nil)
        guard let tiff = image.tiffRepresentation, let png = NSBitmapImageRep(data: tiff)?.representation(using: .png, properties: [:]) else {
            throw ApiError(status: 0, message: "couldn't encode the capture as PNG")
        }
        return png
    }
    #else
    func capture(_ view: WKWebView) async throws -> Data {
        guard view.window != nil, !view.isHidden else {
            throw ApiError(status: 0, message: "the tab isn't on screen: the human has the page view closed or another tab showing")
        }
        let image = try await view.takeSnapshot(configuration: nil)
        guard let png = image.pngData() else {
            throw ApiError(status: 0, message: "couldn't encode the capture as PNG")
        }
        return png
    }
    #endif

    func close() {
        listener?.close()
        listener = nil
        port = nil
    }

    static func fetch(_ api: API, _ taskId: String, _ path: String) async -> (status: Int, type: String, body: Data) {
        do {
            let (body, response) = try await api.session.data(from: api.url("/api/tasks/\(taskId)/raw", [URLQueryItem(name: "path", value: path)]))
            let http = response as! HTTPURLResponse
            return (http.statusCode, http.value(forHTTPHeaderField: "Content-Type") ?? "application/octet-stream", body)
        } catch {
            return (502, "text/plain; charset=utf-8", Data(errorText(error).utf8))
        }
    }
}

private let headEnd = Data("\r\n\r\n".utf8)
private let maxHead = 64 << 10

private final class PageListener: @unchecked Sendable {
    private let api: API
    private let queue = DispatchQueue(label: "run.fc.byre.pages")
    private let served: @Sendable (String, String, Data) -> Void
    private var listener: NWListener?

    init(api: API, ready: @escaping @Sendable (Int) -> Void, served: @escaping @Sendable (String, String, Data) -> Void) {
        self.api = api
        self.served = served
        let params = NWParameters.tcp
        params.requiredLocalEndpoint = .hostPort(host: .ipv4(.loopback), port: .any)
        let listener: NWListener
        do {
            listener = try NWListener(using: params)
        } catch {
            log.error("page listener: \(String(describing: error), privacy: .public)")
            return
        }
        self.listener = listener
        listener.stateUpdateHandler = { [weak listener] state in
            switch state {
            case .ready:
                if let port = listener?.port {
                    ready(Int(port.rawValue))
                }
            case .failed(let error):
                log.error("page listener: \(String(describing: error), privacy: .public)")
            default:
                break
            }
        }
        listener.newConnectionHandler = { [weak self] conn in
            guard let self else {
                conn.cancel()
                return
            }
            conn.start(queue: self.queue)
            self.read(conn, Data())
        }
        listener.start(queue: queue)
    }

    func close() {
        queue.async {
            self.listener?.newConnectionHandler = nil
            self.listener?.cancel()
            self.listener = nil
        }
    }

    private func read(_ conn: NWConnection, _ buf: Data) {
        conn.receive(minimumIncompleteLength: 1, maximumLength: maxHead) { data, _, complete, error in
            var buf = buf
            if let data {
                buf.append(data)
            }
            if let end = buf.range(of: headEnd) {
                self.serve(conn, String(decoding: buf[..<end.lowerBound], as: UTF8.self))
                return
            }
            if complete || error != nil || buf.count > maxHead {
                conn.cancel()
                return
            }
            self.read(conn, buf)
        }
    }

    private func serve(_ conn: NWConnection, _ head: String) {
        let lines = head.components(separatedBy: "\r\n")
        let hosts = lines.dropFirst().compactMap { line -> String? in
            guard let colon = line.firstIndex(of: ":"), line[..<colon].lowercased() == "host" else {
                return nil
            }
            return line[line.index(after: colon)...].trimmingCharacters(in: .whitespaces)
        }
        // A browser must not reach this authenticated file bridge through a
        // hostname that an attacker can rebind to loopback.
        guard let port = listener?.port, hosts == ["127.0.0.1:\(port.rawValue)"] else {
            respond(conn, 403, "text/plain; charset=utf-8", Data("forbidden host".utf8))
            return
        }
        let parts = lines[0].split(separator: " ")
        guard parts.count == 3, parts[0] == "GET", let path = URLComponents(string: String(parts[1]))?.path else {
            respond(conn, 400, "text/plain; charset=utf-8", Data("bad request".utf8))
            return
        }
        let route = path.drop(while: { $0 == "/" }).split(separator: "/", maxSplits: 1, omittingEmptySubsequences: false)
        guard route.count == 2, !route[1].isEmpty else {
            respond(conn, 404, "text/plain; charset=utf-8", Data("not found".utf8))
            return
        }
        let taskId = String(route[0])
        let file = route[1].hasSuffix("/") ? "\(route[1])index.html" : String(route[1])
        let api = api
        Task {
            let (status, type, body) = await Pages.fetch(api, taskId, file)
            self.served(taskId, file, body)
            self.respond(conn, status, type, body)
        }
    }

    private func respond(_ conn: NWConnection, _ status: Int, _ type: String, _ body: Data) {
        let head = "HTTP/1.1 \(status) \(HTTPURLResponse.localizedString(forStatusCode: status))\r\nContent-Type: \(type)\r\nContent-Length: \(body.count)\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n"
        conn.send(content: Data(head.utf8) + body, isComplete: true, completion: .contentProcessed { _ in
            conn.cancel()
        })
    }
}

@MainActor
@Observable
private final class PageWatch {
    private(set) var reloads = 0
    @ObservationIgnored private var timer: Task<Void, Never>?

    func check(_ api: API, _ pages: Pages, _ taskId: String) {
        timer?.cancel()
        timer = Task {
            try? await Task.sleep(for: .milliseconds(1500))
            var changed = false
            for (path, shown) in pages.served[taskId] ?? [:] {
                let fresh = await Pages.fetch(api, taskId, path).body
                if Task.isCancelled {
                    return
                }
                if fresh != shown {
                    pages.served[taskId]?[path] = fresh
                    changed = true
                }
            }
            if changed {
                reloads += 1
            }
        }
    }

    func stop() {
        timer?.cancel()
        timer = nil
    }
}

struct PageTabs: View {
    let shown: ShownPages
    @Environment(Store.self) private var store
    @ViewState private var watch = PageWatch()

    var body: some View {
        VStack(spacing: 0) {
            ScrollView(.horizontal, showsIndicators: false) {
                HStack(spacing: 4) {
                    ForEach(shown.pages, id: \.self) { page in
                        PageTabButton(page: page, active: page == shown.current)
                    }
                }
                .padding(.horizontal, 12)
            }
            .fixedSize(horizontal: false, vertical: true)
            .overlay(alignment: .bottom) {
                Rectangle().fill(Palette.border).frame(height: 1)
            }
            ZStack {
                ForEach(shown.pages, id: \.self) { page in
                    if let view = store.pageView(page) {
                        WebPage(view: view, reloads: page.isFile ? watch.reloads : 0, shown: page == shown.current, refresh: page.isFile)
                    } else if page == shown.current {
                        EmptyText(text: "Loading…")
                    }
                }
            }
        }
        .task(id: shown.current.taskId) {
            let taskId = shown.current.taskId
            let watch = watch
            let api = store.api
            let pages = store.pages
            let off = store.selectedFeed?.appended.add { events in
                if events.contains(where: { $0.kind == "agent.tool_result" || $0.kind == "turn.ended" }) {
                    watch.check(api, pages, taskId)
                }
            }
            await untilCancelled()
            off?()
            watch.stop()
        }
    }
}

private struct PageTabButton: View {
    let page: Page
    let active: Bool
    @Environment(Store.self) private var store
    @ViewState private var hover = false

    var body: some View {
        HStack(spacing: 0) {
            PaneTabButton(label: page.title(store), active: active) {
                store.openPage(page)
            }
            Button {
                store.closePage(page)
            } label: {
                Text("✕")
                    .font(Fonts.font(11))
                    .foregroundStyle(hover ? Palette.text : Palette.muted)
                    .padding(.trailing, 6)
                    .contentShape(Rectangle())
            }
            .buttonStyle(.plain)
            .help("Close this page")
            .onHover { hover = $0 }
        }
        .help(page.url(store)?.absoluteString ?? "")
    }
}

@MainActor
private final class PageDelegate: NSObject, WKUIDelegate, WKNavigationDelegate {
    var popup: ((WKWebView, WKWebViewConfiguration) -> WKWebView?)?
    private var waits: [ObjectIdentifier: CheckedContinuation<Void, any Error>] = [:]

    func wait(_ navigation: WKNavigation?) async throws {
        guard let navigation else {
            return
        }
        try await withCheckedThrowingContinuation { c in
            waits[ObjectIdentifier(navigation)] = c
        }
    }

    private func finish(_ navigation: WKNavigation?, _ error: (any Error)?) {
        guard let navigation, let c = waits.removeValue(forKey: ObjectIdentifier(navigation)) else {
            return
        }
        if let error {
            c.resume(throwing: error)
        } else {
            c.resume()
        }
    }

    func webView(_ webView: WKWebView, createWebViewWith configuration: WKWebViewConfiguration, for action: WKNavigationAction, windowFeatures: WKWindowFeatures) -> WKWebView? {
        popup?(webView, configuration)
    }

    func webView(_ webView: WKWebView, runJavaScriptAlertPanelWithMessage message: String, initiatedByFrame frame: WKFrameInfo) async {
        _ = await ask(webView, frame, message, cancel: false, field: nil)
    }

    func webView(_ webView: WKWebView, runJavaScriptConfirmPanelWithMessage message: String, initiatedByFrame frame: WKFrameInfo) async -> Bool {
        await ask(webView, frame, message, cancel: true, field: nil) != nil
    }

    func webView(_ webView: WKWebView, runJavaScriptTextInputPanelWithPrompt prompt: String, defaultText: String?, initiatedByFrame frame: WKFrameInfo) async -> String? {
        await ask(webView, frame, prompt, cancel: true, field: defaultText ?? "")
    }

    // Returns nil when the human cancels, otherwise the field's text ("" without a field).
    private func ask(_ webView: WKWebView, _ frame: WKFrameInfo, _ message: String, cancel: Bool, field: String?) async -> String? {
        let title = frame.securityOrigin.host.isEmpty ? "This page says" : "\(frame.securityOrigin.host) says"
        #if os(macOS)
        guard let window = unsafe webView.window, !webView.isHiddenOrHasHiddenAncestor else {
            return nil
        }
        let alert = NSAlert()
        alert.messageText = title
        alert.informativeText = message
        alert.addButton(withTitle: "OK")
        if cancel {
            alert.addButton(withTitle: "Cancel")
        }
        let input = field.map { text in
            let f = NSTextField(frame: NSRect(x: 0, y: 0, width: 300, height: 24))
            f.stringValue = text
            return f
        }
        alert.accessoryView = input
        guard await alert.beginSheetModal(for: window) == .alertFirstButtonReturn else {
            return nil
        }
        return input?.stringValue ?? ""
        #else
        guard let window = webView.window, !webView.isHidden, var top = window.rootViewController else {
            return nil
        }
        while let presented = top.presentedViewController {
            top = presented
        }
        let presenter = top
        return await withCheckedContinuation { c in
            let alert = UIAlertController(title: title, message: message, preferredStyle: .alert)
            if let field {
                alert.addTextField { f in
                    f.text = field
                }
            }
            alert.addAction(UIAlertAction(title: "OK", style: .default) { [weak alert] _ in
                c.resume(returning: alert?.textFields?.first?.text ?? "")
            })
            if cancel {
                alert.addAction(UIAlertAction(title: "Cancel", style: .cancel) { _ in
                    c.resume(returning: nil)
                })
            }
            presenter.present(alert, animated: true)
        }
        #endif
    }

    func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
        finish(navigation, nil)
    }

    func webView(_ webView: WKWebView, didFail navigation: WKNavigation!, withError error: any Error) {
        finish(navigation, error)
    }

    func webView(_ webView: WKWebView, didFailProvisionalNavigation navigation: WKNavigation!, withError error: any Error) {
        finish(navigation, error)
    }
}

@MainActor
private struct WebPage {
    let view: WKWebView
    let reloads: Int
    let shown: Bool
    let refresh: Bool

    final class Coordinator {
        var reloads = 0
    }

    func makeCoordinator() -> Coordinator {
        Coordinator()
    }

    private func make(_ c: Coordinator) -> WKWebView {
        c.reloads = reloads
        if refresh, !view.isLoading {
            view.reload()
        }
        return view
    }

    private func update(_ view: WKWebView, _ c: Coordinator) {
        view.isHidden = !shown
        if c.reloads != reloads {
            c.reloads = reloads
            view.reload()
        }
    }
}

#if os(macOS)
extension WebPage: NSViewRepresentable {
    func makeNSView(context: Context) -> WKWebView {
        make(context.coordinator)
    }

    func updateNSView(_ view: WKWebView, context: Context) {
        update(view, context.coordinator)
    }
}
#else
extension WebPage: UIViewRepresentable {
    func makeUIView(context: Context) -> WKWebView {
        make(context.coordinator)
    }

    func updateUIView(_ view: WKWebView, context: Context) {
        update(view, context.coordinator)
    }
}
#endif