import Combine
import ComposableArchitecture
import XCTest

final class ObservableTests: BaseTCATestCase {
  func testBasics() async {
    var state = ChildState()
    let countDidChange = self.expectation(description: "count.didChange")

    withPerceptionTracking {
      _ = state.count
    } onChange: {
      countDidChange.fulfill()
    }

    state.count += 1
    await self.fulfillment(of: [countDidChange], timeout: 0)
    XCTAssertEqual(state.count, 1)
  }

  func testAssignEqualValue() async {
    var state = ChildState()
    let didChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.count
    } onChange: {
      didChange.withValue { $0 = true }
    }

    state.count = state.count
    XCTAssertEqual(state.count, 0)
    XCTAssert(!didChange.withValue { $0 })
  }

  func testCopyMutation() async {
    XCTTODO(
      """
      Ideally this test would pass but it does not because making a copy of a child state, mutating
      it, and assigning it does not change the identified array's IDs, and therefore the fast-path
      of _$isIdentityEqual prevents observation.
      """
    )

    var state = ParentState(children: [ChildState(count: 42)])
    let countDidChange = self.expectation(description: "count.didChange")
    var copy = state.children[0]
    copy.count += 1

    withPerceptionTracking {
      _ = state.children[0].count
    } onChange: {
      countDidChange.fulfill()
    }

    state.children[0] = copy

    await self.fulfillment(of: [countDidChange], timeout: 0.1)
    XCTAssertEqual(state.children[0].count, 43)
  }

  func testCopyMutation_WithPerturbation() async {
    var state = ParentState(children: [ChildState(count: 42)])
    let countDidChange = self.expectation(description: "count.didChange")
    var copy = state.children[0]
    copy.count += 1

    withPerceptionTracking {
      _ = state.children[0].count
    } onChange: {
      countDidChange.fulfill()
    }

    state.children[0] = copy
    state.children[0]._$willModify()

    await self.fulfillment(of: [countDidChange], timeout: 0.1)
    XCTAssertEqual(state.children[0].count, 43)
  }

  func testReplace() async {
    #if swift(<6.2)
      if #available(iOS 17, macOS 14, tvOS 14, watchOS 10, *) {
        XCTTODO("Ideally this would pass but we cannot detect this kind of mutation currently.")
      }
    #endif
    var state = ChildState(count: 42)
    let didChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.count
    } onChange: {
      didChange.withValue { $0 = true }
    }

    state.replace(with: ChildState())
    XCTAssertEqual(state.count, 0)
    XCTAssert(didChange.withValue { $0 })
  }

  func testReset() async {
    #if swift(<6.2)
      if #available(iOS 17, macOS 14, tvOS 14, watchOS 10, *) {
        XCTTODO("Ideally this would pass but we cannot detect this kind of mutation currently.")
      }
    #endif

    var state = ChildState(count: 42)
    let didChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.count
    } onChange: {
      didChange.withValue { $0 = true }
    }

    state.reset()
    XCTAssertEqual(state.count, 0)
    XCTAssert(didChange.withValue { $0 })
  }

  func testChildCountMutation() async {
    var state = ParentState()
    let childCountDidChange = LockIsolated(false)
    let childDidChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.child.count
    } onChange: {
      childCountDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.child
    } onChange: {
      childDidChange.withValue { $0 = true }
    }

    state.child.count += 1
    XCTAssertEqual(state.child.count, 1)
    XCTAssert(childCountDidChange.withValue { $0 })
    XCTAssert(!childDidChange.withValue { $0 })
  }

  func testChildReset() async {
    var state = ParentState()
    let childCountDidChange = LockIsolated(false)
    let childDidChange = LockIsolated(false)

    let child = state.child
    withPerceptionTracking {
      _ = child.count
    } onChange: {
      childCountDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.child
    } onChange: {
      childDidChange.withValue { $0 = true }
    }

    state.child = ChildState(count: 42)
    XCTAssertEqual(state.child.count, 42)
    XCTAssert(!childCountDidChange.withValue { $0 })
    XCTAssert(childDidChange.withValue { $0 })
  }

  func testReplaceChild() async {
    var state = ParentState()
    let childDidChange = self.expectation(description: "child.didChange")

    withPerceptionTracking {
      _ = state.child
    } onChange: {
      childDidChange.fulfill()
    }

    state.child.replace(with: ChildState(count: 42))
    await self.fulfillment(of: [childDidChange], timeout: 0)
    XCTAssertEqual(state.child.count, 42)
  }

  func testResetChild() async {
    var state = ParentState(child: ChildState(count: 42))
    let childDidChange = self.expectation(description: "child.didChange")

    withPerceptionTracking {
      _ = state.child
    } onChange: {
      childDidChange.fulfill()
    }

    state.child.reset()
    await self.fulfillment(of: [childDidChange], timeout: 0)
    XCTAssertEqual(state.child.count, 0)
  }

  func testSwapSiblings() async {
    var state = ParentState(
      child: ChildState(count: 1),
      sibling: ChildState(count: -1)
    )
    let childDidChange = self.expectation(description: "child.didChange")
    let siblingDidChange = self.expectation(description: "sibling.didChange")

    withPerceptionTracking {
      _ = state.child
    } onChange: {
      childDidChange.fulfill()
    }
    withPerceptionTracking {
      _ = state.sibling
    } onChange: {
      siblingDidChange.fulfill()
    }

    state.swap()
    await self.fulfillment(of: [childDidChange], timeout: 0)
    await self.fulfillment(of: [siblingDidChange], timeout: 0)
    XCTAssertEqual(state.child.count, -1)
    XCTAssertEqual(state.sibling.count, 1)
  }

  func testPresentOptional() async {
    var state = ParentState()
    let optionalDidChange = self.expectation(description: "optional.didChange")

    withPerceptionTracking {
      _ = state.optional
    } onChange: {
      optionalDidChange.fulfill()
    }

    state.optional = ChildState(count: 42)
    await self.fulfillment(of: [optionalDidChange], timeout: 0)
    XCTAssertEqual(state.optional?.count, 42)
  }

  func testMutatePresentedOptional() async {
    var state = ParentState(optional: ChildState())
    let optionalDidChange = LockIsolated(false)
    let optionalCountDidChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.optional
    } onChange: {
      optionalDidChange.withValue { $0 = true }
    }
    let optional = state.optional
    withPerceptionTracking {
      _ = optional?.count
    } onChange: {
      optionalCountDidChange.withValue { $0 = true }
    }

    state.optional?.count += 1
    XCTAssertEqual(state.optional?.count, 1)
    XCTAssert(!optionalDidChange.withValue { $0 })
    XCTAssert(optionalCountDidChange.withValue { $0 })
  }

  func testPresentDestination() async {
    var state = ParentState()
    let destinationDidChange = self.expectation(description: "destination.didChange")

    withPerceptionTracking {
      _ = state.destination
    } onChange: {
      destinationDidChange.fulfill()
    }

    state.destination = .child1(ChildState(count: 42))
    await self.fulfillment(of: [destinationDidChange], timeout: 0)
    XCTAssertEqual(state.destination?.child1?.count, 42)
  }

  func testDismissDestination() async {
    var state = ParentState(destination: .child1(ChildState()))
    let destinationDidChange = self.expectation(description: "destination.didChange")

    withPerceptionTracking {
      _ = state.destination
    } onChange: {
      destinationDidChange.fulfill()
    }

    state.destination = nil
    await self.fulfillment(of: [destinationDidChange], timeout: 0)
    XCTAssertEqual(state.destination, nil)
  }

  func testChangeDestination() async {
    var state = ParentState(destination: .child1(ChildState()))
    let destinationDidChange = self.expectation(description: "destination.didChange")

    withPerceptionTracking {
      _ = state.destination
    } onChange: {
      destinationDidChange.fulfill()
    }

    state.destination = .child2(ChildState(count: 42))
    await self.fulfillment(of: [destinationDidChange], timeout: 0)
    XCTAssertEqual(state.destination?.child2?.count, 42)
  }

  func testChangeDestination_KeepIdentity() async {
    let childState = ChildState(count: 42)
    var state = ParentState(destination: .child1(childState))
    let destinationDidChange = self.expectation(description: "destination.didChange")

    withPerceptionTracking {
      _ = state.destination
    } onChange: {
      destinationDidChange.fulfill()
    }

    state.destination = .child2(childState)
    await self.fulfillment(of: [destinationDidChange], timeout: 0)
    XCTAssertEqual(state.destination?.child2?.count, 42)
  }

  func testMutatingDestination_NonObservableCase() async {
    let expectation = self.expectation(description: "destination.didChange")
    var state = ParentState(destination: .inert(0))

    withPerceptionTracking {
      _ = state.destination
    } onChange: {
      expectation.fulfill()
    }

    state.destination = .inert(1)
    XCTAssertEqual(state.destination, .inert(1))
    await self.fulfillment(of: [expectation])
  }

  func testReplaceWithCopy() async {
    let childState = ChildState(count: 1)
    var childStateCopy = childState
    childStateCopy.count = 2
    var state = ParentState(child: childState, sibling: childStateCopy)
    let childCountDidChange = self.expectation(description: "child.count.didChange")

    withPerceptionTracking {
      _ = state.child.count
    } onChange: {
      childCountDidChange.fulfill()
    }

    state.child.replace(with: state.sibling)

    await self.fulfillment(of: [childCountDidChange], timeout: 0)
    XCTAssertEqual(state.child.count, 2)
    XCTAssertEqual(state.sibling.count, 2)
  }

  @MainActor
  func testStore_ReplaceChild() async {
    let store = Store<ParentState, Void>(initialState: ParentState()) {
      Reduce { state, _ in
        state.child.replace(with: ChildState(count: 42))
        return .none
      }
    }
    let childDidChange = self.expectation(description: "child.didChange")

    withPerceptionTracking {
      _ = store.child
    } onChange: {
      childDidChange.fulfill()
    }

    store.send(())
    await self.fulfillment(of: [childDidChange], timeout: 0)
    XCTAssertEqual(store.child.count, 42)
  }

  @MainActor
  func testStore_Replace() async {
    let store = Store<ChildState, Void>(initialState: ChildState()) {
      Reduce { state, _ in
        state.replace(with: ChildState(count: 42))
        return .none
      }
    }
    let countDidChange = self.expectation(description: "child.didChange")

    withPerceptionTracking {
      _ = store.count
    } onChange: {
      countDidChange.fulfill()
    }

    store.send(())
    await self.fulfillment(of: [countDidChange], timeout: 0)
    XCTAssertEqual(store.count, 42)
  }

  @MainActor
  func testStore_ResetChild() async {
    let store = Store<ParentState, Void>(initialState: ParentState(child: ChildState(count: 42))) {
      Reduce { state, _ in
        state.child.reset()
        return .none
      }
    }
    let childDidChange = self.expectation(description: "child.didChange")

    withPerceptionTracking {
      _ = store.child
    } onChange: {
      childDidChange.fulfill()
    }

    store.send(())
    await self.fulfillment(of: [childDidChange], timeout: 0)
    XCTAssertEqual(store.child.count, 0)
  }

  @MainActor
  func testStore_Reset() async {
    let store = Store<ChildState, Void>(initialState: ChildState(count: 42)) {
      Reduce { state, _ in
        state.reset()
        return .none
      }
    }
    let countDidChange = self.expectation(description: "child.didChange")

    withPerceptionTracking {
      _ = store.count
    } onChange: {
      countDidChange.fulfill()
    }

    store.send(())
    await self.fulfillment(of: [countDidChange], timeout: 0)
    XCTAssertEqual(store.count, 0)
  }

  func testIdentifiedArray_AddElement() {
    var state = ParentState()
    let rowsDidChange = self.expectation(description: "rowsDidChange")

    withPerceptionTracking {
      _ = state.rows
    } onChange: {
      rowsDidChange.fulfill()
    }

    state.rows.append(ChildState())
    XCTAssertEqual(state.rows.count, 1)
    self.wait(for: [rowsDidChange], timeout: 0)
  }

  func testIdentifiedArray_MutateElement() {
    var state = ParentState(rows: [
      ChildState(),
      ChildState(),
    ])
    let rowsDidChange = LockIsolated(false)
    let firstRowDidChange = LockIsolated(false)
    let firstRowCountDidChange = LockIsolated(false)
    let secondRowDidCountChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.rows
    } onChange: {
      rowsDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.rows[0]
    } onChange: {
      firstRowDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.rows[0].count
    } onChange: {
      firstRowCountDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.rows[1].count
    } onChange: {
      secondRowDidCountChange.withValue { $0 = true }
    }

    state.rows[0].count += 1
    XCTAssertEqual(state.rows[0].count, 1)
    XCTAssert(!rowsDidChange.withValue { $0 })
    XCTAssert(!firstRowDidChange.withValue { $0 })
    XCTAssert(firstRowCountDidChange.withValue { $0 })
    XCTAssert(!secondRowDidCountChange.withValue { $0 })
  }

  func testIdentifiedArray_CopyMutateAssignElement() {
    var state = ParentState(rows: [
      ChildState(),
      ChildState(),
    ])
    let rowsDidChange = LockIsolated(0)
    let firstRowDidChange = LockIsolated(0)
    let firstRowCountDidChange = LockIsolated(0)
    let secondRowDidCountChange = LockIsolated(0)

    withPerceptionTracking {
      _ = state.rows
    } onChange: {
      rowsDidChange.withValue { $0 += 1 }
    }
    withPerceptionTracking {
      _ = state.rows[0]
    } onChange: {
      firstRowDidChange.withValue { $0 += 1 }
    }
    withPerceptionTracking {
      _ = state.rows[0].count
    } onChange: {
      firstRowCountDidChange.withValue { $0 += 1 }
    }
    withPerceptionTracking {
      _ = state.rows[1].count
    } onChange: {
      secondRowDidCountChange.withValue { $0 += 1 }
    }

    var copy = state.rows[0]
    copy.count += 1
    state.rows[id: copy.id] = copy

    XCTAssertEqual(state.rows[0].count, 1)
    XCTAssert(rowsDidChange.withValue(\.self) == 1)
    XCTAssert(firstRowDidChange.withValue(\.self) == 1)
    XCTAssert(firstRowCountDidChange.withValue(\.self) == 1)
    XCTAssert(secondRowDidCountChange.withValue(\.self) == 1)

    state.rows[id: state.rows[0].id] = nil
    XCTAssertEqual(state.rows.count, 1)
    XCTAssertEqual(state.rows[0].count, 0)
    XCTAssertEqual(rowsDidChange.withValue(\.self), 2)
    XCTAssertEqual(firstRowDidChange.withValue(\.self), 1)
    XCTAssertEqual(firstRowCountDidChange.withValue(\.self), 1)
    XCTAssertEqual(secondRowDidCountChange.withValue(\.self), 1)

    let new = ChildState(count: 42)
    state.rows[id: new.id] = new
    XCTAssertEqual(state.rows[1].count, 42)
    XCTAssertEqual(rowsDidChange.withValue(\.self), 3)
    XCTAssertEqual(firstRowDidChange.withValue(\.self), 1)
    XCTAssertEqual(firstRowCountDidChange.withValue(\.self), 1)
    XCTAssertEqual(secondRowDidCountChange.withValue(\.self), 1)
  }

  func testPresents_NilToNonNil() {
    var state = ParentState()
    let presentationDidChange = self.expectation(description: "presentationDidChange")

    withPerceptionTracking {
      _ = state.presentation
    } onChange: {
      presentationDidChange.fulfill()
    }

    state.presentation = ChildState()
    XCTAssertEqual(state.presentation?.count, 0)
    self.wait(for: [presentationDidChange], timeout: 0)
  }

  func testPresents_Mutate() {
    var state = ParentState(presentation: ChildState())
    let presentationDidChange = LockIsolated(false)
    let presentationCountDidChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.presentation
    } onChange: {
      presentationDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.presentation?.count
    } onChange: {
      presentationCountDidChange.withValue { $0 = true }
    }

    state.presentation?.count += 1
    XCTAssertEqual(state.presentation?.count, 1)
    XCTAssert(!presentationDidChange.withValue { $0 })
    XCTAssert(presentationCountDidChange.withValue { $0 })
  }

  func testStackState_AddElement() {
    var state = ParentState()
    let pathDidChange = self.expectation(description: "pathDidChange")

    withPerceptionTracking {
      _ = state.path
    } onChange: {
      pathDidChange.fulfill()
    }

    state.path.append(ChildState())
    XCTAssertEqual(state.path.count, 1)
    self.wait(for: [pathDidChange], timeout: 0)
  }

  func testStackState_MutateElement() {
    var state = ParentState(
      path: StackState([
        ChildState(),
        ChildState(),
      ])
    )
    let pathDidChange = LockIsolated(false)
    let firstElementDidChange = LockIsolated(false)
    let firstElementCountDidChange = LockIsolated(false)
    let secondElementCountDidChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.path
    } onChange: {
      pathDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.path[0]
    } onChange: {
      firstElementDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.path[0].count
    } onChange: {
      firstElementCountDidChange.withValue { $0 = true }
    }
    withPerceptionTracking {
      _ = state.path[1].count
    } onChange: {
      secondElementCountDidChange.withValue { $0 = true }
    }

    state.path[id: 0]?.count += 1
    XCTAssertEqual(state.path[0].count, 1)
    XCTAssert(!pathDidChange.withValue { $0 })
    XCTAssert(!firstElementDidChange.withValue { $0 })
    XCTAssert(firstElementCountDidChange.withValue { $0 })
    XCTAssert(!secondElementCountDidChange.withValue { $0 })
  }

  func testCopy() {
    var state = ParentState()
    var childCopy = state.child.copy()
    childCopy.count = 42
    let childCountDidChange = self.expectation(description: "childCountDidChange")

    withPerceptionTracking {
      _ = state.child.count
    } onChange: {
      childCountDidChange.fulfill()
    }

    state.child.replace(with: childCopy)
    XCTAssertEqual(state.child.count, 42)
    self.wait(for: [childCountDidChange], timeout: 0)
  }

  func testArrayAppend() {
    var state = ParentState()
    let childrenDidChange = self.expectation(description: "childrenDidChange")

    withPerceptionTracking {
      _ = state.children
    } onChange: {
      childrenDidChange.fulfill()
    }

    state.children.append(ChildState())
    self.wait(for: [childrenDidChange])
  }

  func testArrayMutate() {
    var state = ParentState(children: [ChildState()])
    var didChange = LockIsolated(false)

    withPerceptionTracking {
      _ = state.children
    } onChange: {
      didChange.withValue { $0 = true }
    }

    state.children[0].count += 1
    XCTAssert(!didChange.withValue { $0 })
  }

  @MainActor
  func testEnumStateWithInertCases() {
    let store = Store<EnumState, Void>(initialState: EnumState.count(.one)) {
      Reduce { state, _ in
        state = .count(.two)
        return .none
      }
    }
    let onChangeExpectation = self.expectation(description: "onChange")
    withPerceptionTracking {
      _ = store.state
    } onChange: {
      onChangeExpectation.fulfill()
    }

    store.send(())

    self.wait(for: [onChangeExpectation], timeout: 0)
  }

  @MainActor
  func testEnumStateWithInertCasesTricky() {
    let store = Store<EnumState, Void>(initialState: EnumState.count(.one)) {
      Reduce { state, _ in
        state = .anotherCount(.one)
        return .none
      }
    }
    let onChangeExpectation = self.expectation(description: "onChange")
    withPerceptionTracking {
      _ = store.state
    } onChange: {
      onChangeExpectation.fulfill()
    }

    store.send(())

    self.wait(for: [onChangeExpectation], timeout: 0)
  }

  @MainActor
  func testEnumStateWithIntCase() {
    let store = Store<EnumState, Void>(initialState: EnumState.int(0)) {
      Reduce { state, _ in
        state = .int(1)
        return .none
      }
    }
    let onChangeExpectation = self.expectation(description: "onChange")
    withPerceptionTracking {
      _ = store.state
    } onChange: {
      onChangeExpectation.fulfill()
    }

    store.send(())

    self.wait(for: [onChangeExpectation], timeout: 0)
  }
}

@ObservableState
private struct ChildState: Equatable, Identifiable {
  var count = 0
  mutating func replace(with other: Self) {
    self = other
  }
  mutating func reset() {
    self = Self()
  }
  mutating func copy() -> Self {
    self
  }
}
@ObservableState
private struct ParentState: Equatable {
  var child = ChildState()
  @Presents var destination: DestinationState?
  var children: [ChildState] = []
  @Presents var optional: ChildState?
  var path = StackState<ChildState>()
  @Presents var presentation: ChildState?
  var rows: IdentifiedArrayOf<ChildState> = []
  var sibling = ChildState()
  mutating func swap() {
    let childCopy = child
    self.child = self.sibling
    self.sibling = childCopy
  }
}
@dynamicMemberLookup
@CasePathable
@ObservableState
private enum DestinationState: Equatable {
  case child1(ChildState)
  case child2(ChildState)
  case inert(Int)
}
@ObservableState
private enum EnumState: Equatable {
  case count(Count)
  case anotherCount(Count)
  case int(Int)
  @ObservableState
  enum Count: String {
    case one, two
  }
}