dependencies {
    implementation("io.github.brianmwas.koog_compose:koog-compose-core:2.1.0")

    // Optional modules — add what you need
    implementation("io.github.brianmwas.koog_compose:koog-compose-ui:2.1.0")           // Material 3 chat components
    implementation("io.github.brianmwas.koog_compose:koog-compose-device:2.1.0")       // GPS, alarms, WorkManager (Android)
    implementation("io.github.brianmwas.koog_compose:koog-compose-mediapipe:2.1.0")    // On-device models (Gemma, Apple FMs)
    implementation("io.github.brianmwas.koog_compose:koog-compose-session-room:2.1.0") // Room-backed session persistence
    implementation("io.github.brianmwas.koog_compose:koog-compose-testing:2.1.0")      // Test utilities
}

import io.github.koogcompose.provider.ondevice.installOnDeviceProviderSupport

fun initAi() {
    installOnDeviceProviderSupport()  // Application.onCreate() or main()
}

@Serializable
data class RunState(
    val userName: String,
    val isRunning: Boolean = false,
    val distanceKm: Double = 0.0,
    val durationMs: Long = 0,
    val pace: String? = null,
)

val runCoach = koogCompose<RunState> {
    provider {
        onDevice(modelPath = "/data/models/gemma-4-E2B.litertlm") {
            onUnavailable { anthropic(apiKey = BuildConfig.KEY) }
        }
    }

    initialState { RunState(userName = "brian") }

    phases {
        phase("ready", initial = true) {
            instructions { "Ask the user if they're ready for their run." }
            tool(StartRunTimerTool(stateStore))
        }

        phase("running") {
            instructions { "The run is active. Check in if they go quiet for 15 minutes." }
            tool(BackgroundTimerTool())         // WorkManager — survives the app closing
            tool(LocationTrackerTool(stateStore))
        }

        phase("finished") {
            instructions { "Summarise their run: duration, distance, pace." }
            tool(StopTimerTool(stateStore))
            tool(CalculatePaceTool(stateStore))
        }
    }

    config {
        retry { maxAttempts = 3; initialDelayMs = 500L }
        stuckDetection { threshold = 3; fallbackMessage = "Let me try a different approach." }
    }
}

@Composable
fun RunScreen(definition: KoogDefinition<RunState> = koogCompose { ... }) {
    val session = rememberPhaseSession(definition) {
        sessionId = "run_brian"
    }
    val runState by session.appState.collectAsState()

    if (runState.isRunning) {
        Text("Running — ${runState.distanceKm} km")
    }

    Scaffold(bottomBar = { ChatInputBar(rememberChatState(session)) }) { padding ->
        ChatMessageList(rememberChatState(session), modifier = Modifier.padding(padding))
    }
}

class RunViewModel(context: KoogComposeContext<RunState>, executor: PromptExecutor) : ViewModel() {
    val session = phaseSession(context, executor) {
        sessionId = "run_brian"
        scope = viewModelScope
    }

    val responseStream = session.responseStream  // Flow<String> — tokens as they arrive
    val runState       = session.appState        // StateFlow<RunState>
}

phase("finish_run") {
    subphase("stop_timer") {
        instructions { "Stop the run timer and record final duration." }
        tool(StopTimerTool(stateStore))
    }
    subphase("calculate_stats") {
        instructions { "Calculate distance and pace from the GPS trace." }
        tool(CalculatePaceTool(stateStore))
    }
    subphase("save_run") {
        instructions { "Save the run to storage." }
        tool(SaveRunTool(stateStore))
    }
    onCondition("run saved", "summary")
}

phase("gather_context", initial = true) {
    parallel {
        branch("weather")  { tool(WeatherTool(stateStore)) }
        branch("location") { tool(GeocoderTool(stateStore)) }
        branch("history")  { tool(RunHistoryTool(stateStore)) }
    }
    onCondition("context ready", "plan")
}

class LocationTrackerTool(
    override val stateStore: KoogStateStore<RunState>
) : StatefulTool<RunState>() {
    override val name            = "TrackLocation"
    override val description     = "Record GPS coordinates during the run"
    override val permissionLevel = PermissionLevel.SENSITIVE

    override suspend fun execute(args: JsonObject): ToolResult {
        val location = getCurrentLocation()
        stateStore.update {
            it.copy(gpsTrace = it.gpsTrace + location)
        }
        return ToolResult.Success("Recorded ${location.latitude}, ${location.longitude}")
    }
}

val displayText by remember {
    viewModel.turnId.flatMapLatest { _ ->
        viewModel.responseStream.runningFold("") { acc, token -> acc + token }
    }
}.collectAsState(initial = "")

processor.outcomes.collect { outcome ->
    if (outcome is DirectiveOutcome.Accepted && outcome.positionFallback) {
        // The Before/After reference was missing — component went to the end.
        // The agent can re-issue a ReorderComponents directive if ordering matters.
    }
}

provider {
    onDevice(modelPath = "/data/models/gemma-4-E2B.litertlm") {
        maxToolRounds(8)
        onUnavailable {
            // Fallback only if model is unavailable:
            // - File missing or corrupted
            // - Device hardware incompatible
            // ⚠️ This fallback sends data to Anthropic's servers
            anthropic(apiKey = BuildConfig.KEY)
        }
    }
}

config {
    auditLog { redactArgs = true }
}

val focusAgent = koogAgent("focus") {
    instructions { "You are a focus session specialist." }
    phases { phase("active") { /* ... */ } }
}

val session = koogSession<Unit> {
    provider { ollama(model = "llama3.2") }
    main {
        phases {
            phase("root", initial = true) {
                handoff(focusAgent) {
                    "User asks about focus, productivity, or pomodoro"
                }
            }
        }
    }
    agents(focusAgent)
}

config {
    eventSink = PrintlnEventSink          // dev: logs to console
    // or
    eventSink = FirebaseEventSink()       // prod: Firebase Analytics
    // or
    eventSink = NoOpEventSink             // tests: silent
}

class FirebaseEventSink(private val analytics: FirebaseAnalytics) : EventSink {
    override suspend fun emit(event: AgentEvent) {
        val bundle = when (event) {
            is AgentEvent.ToolCalled -> Bundle().apply {
                putString("toolName", event.toolName)
            }
            is AgentEvent.PhaseTransitioned -> Bundle().apply {
                putString("from", event.from)
                putString("to", event.to)
            }
            else -> Bundle()
        }
        analytics.logEvent(event::class.simpleName ?: "AgentEvent", bundle)
    }
}

// From a notification
session.resumeAt("notify_user", userMessage = "Your run is ready to view!")

// From a deep link — no user message, no history pollution
session.resumeAt("onboarding_flow")

val researchSubphase = subphaseTemplate("research") {
    instructions { "Search and summarise relevant information." }
    tool(WebSearchTool(stateStore))
}

phase("respond") {
    include(researchSubphase)    // adds the "research" subphase
    subphase("compose_answer") { /* ... */ }
}

val session = phaseSession(context, executor) {
    sessionId = "run_brian"
    scope = viewModelScope
    store = RoomSessionStore(db.sessionDao())  // ← Room backend
}

val migration = object : StateMigration<AppState> {
    override val schemaVersion = 3
    override suspend fun migrate(json: JsonObject, fromVersion: Int): JsonObject {
        return when (fromVersion) {
            // v1 → v2: add themeMode field
            1    -> json + ("themeMode" to JsonPrimitive("System"))
            // v2 → v3: rename "userName" → "userDisplayName"
            2    -> (json.toMutableMap() as MutableMap<String, JsonElement>).apply {
                val userName = remove("userName")
                if (userName != null) put("userDisplayName", userName)
            }.let { JsonObject(it) }
            else -> json
        }
    }
}

@Composable
fun MyScreen(definition: KoogDefinition<MyState> = koogCompose { ... }) {
    val session = rememberPhaseSession(definition) {
        sessionId = "my_screen_session"
    }
}

class MyViewModel(context: KoogComposeContext<MyState>, executor: PromptExecutor) : ViewModel() {
    val session = phaseSession(context, executor) {
        sessionId = "my_session"
        scope = viewModelScope
        store = RoomSessionStore(db.sessionDao())  // optional
    }
}

val definition = koogCompose<MyState> { ... }
val session = definition.createPhaseSession(executor, viewModelScope) {
    sessionId = "my_session"
}

class SavePhotoTool : StatefulTool<AppState>() {
    override suspend fun execute(args: JsonObject): ToolResult {
        return try {
            saveFile(args["path"]?.content ?: "")
            ToolResult.Success("Saved")
        } catch (e: IOException) when {
            e.isNetworkRelated() -> ToolResult.Failure(
                message = "Network hiccup. Retrying shortly...",
                retryable = true,                           // Agent can retry automatically
                recoveryHint = RecoveryHint.RetryAfterDelay // With backoff
            )
            e.isStorageFull() -> ToolResult.Denied(
                reason = "Storage full",
                recoveryHint = RecoveryHint.RequiresUserAction(
                    "Please free up space and say 'try again'"
                )
            )
            else -> ToolResult.Failure("Couldn't save", retryable = false)
        }
    }
}

val breaker = CircuitBreaker(failureThreshold = 5, cooldownMs = 60_000)
val tool = CircuitBreakerGuard(
    delegate = SavePhotoTool(stateStore),
    circuitBreaker = breaker,
    sessionId = session.id,        // optional
    eventSink = myEventSink,       // optional — emits CircuitBreakerOpened/Closed
)

val result = store.loadOrRecover(sessionId)
when (result) {
    is SessionLoadResult.Success   -> session = resumeSession(result.session)
    is SessionLoadResult.Recovered -> {
        showMessage(result.reason)  // "Session corrupted, starting fresh"
        session = startNewSession()
    }
    is SessionLoadResult.NotFound  -> { }
}

config {
    retry {
        maxAttempts = 3
        initialDelayMs = 1_000
        backoffMultiplier = 2.0   // 1s → 2s → 4s
    }
}

catch (e: IOException) {
    val userMessage = when {
        e.isNetworkRelated() -> "Internet connection problem — trying again..."
        e.isStorageFull()    -> "Your device is full — please free up space"
        else                 -> "Something went wrong — our team is aware"
    }
    ToolResult.Failure(userMessage, retryable = false)
}

@Test
fun `"I'm back" transitions from running to finished`() {
    val session = testPhaseSession(context) {
        on("I'm back", phase = "running") {
            transitionTo("finished")
            callTool("StopTimer")
            callTool("CalculateStats")
            respondWith("Great run! 3.2 km in 18 minutes — 5:38 pace.")
        }
    }

    session.send("I'm back")

    assertPhase(session, "finished")
    assertToolCalled(session, "StopTimer")
    assertState(session) { assertFalse(it.isRunning) }
}