package com.correx.core.tasks

import com.correx.core.events.types.TaskId
import com.correx.core.events.types.TaskLinkType
import com.correx.core.events.types.TaskTargetKind

/**
 * Assembles a [TaskContextBundle] for a task. Links are partitioned by their recorded
 * [com.correx.core.events.types.TaskTargetKind]: TASK targets resolve to [RelatedTask] dependencies
 * (with live status), DOC targets inline via [documentResolver], ARTIFACT targets via
 * [artifactResolver], SESSION targets via [sessionResolver]. Any kind whose resolver is absent or
 * comes up empty falls back to a raw [RelatedLink].
 *
 * When a [TaskKnowledgeRetriever] is supplied, the bundle is additionally enriched with
 * semantically-relevant snippets over the task's title+goal. Retrieval and resolution failures are
 * swallowed (the bundle still returns) so context assembly never fails on a flaky dependency.
 */
class TaskContextAssembler(
    private val service: TaskService,
    private val retriever: TaskKnowledgeRetriever? = null,
    private val documentResolver: TaskDocumentResolver? = null,
    private val artifactResolver: TaskArtifactResolver? = null,
    private val sessionResolver: TaskSessionResolver? = null,
    private val knowledgeLimit: Int = DEFAULT_KNOWLEDGE_LIMIT,
) {

    suspend fun assemble(taskId: TaskId): TaskContextBundle? {
        val state = service.getTask(taskId)?.state ?: return null

        val resolved = ResolvedLinks()
        for (link in state.links) {
            // Resolution dispatches on the link's recorded targetKind — no guessing from the id.
            // A kind whose resolution comes up empty falls back to a raw related link.
            when (link.targetKind) {
                TaskTargetKind.TASK -> addTask(link, resolved)
                TaskTargetKind.DOC -> addDocument(link, resolved)
                TaskTargetKind.ARTIFACT -> addArtifact(link, resolved)
                TaskTargetKind.SESSION -> addSession(link, resolved)
            }
        }

        val blockedBy = computeBlockedBy(taskId, state)

        return TaskContextBundle(
            id = taskId.value,
            key = state.key,
            status = state.status.name,
            title = state.title,
            goal = state.goal,
            acceptanceCriteria = state.acceptanceCriteria,
            relevantFiles = state.affectedPaths,
            blocked = blockedBy.isNotEmpty(),
            blockedBy = blockedBy,
            dependencies = resolved.dependencies,
            documents = resolved.documents,
            artifacts = resolved.artifacts,
            sessions = resolved.sessions,
            relatedLinks = resolved.related,
            notes = state.notes.map { NoteEntry(it.author.name, it.body) },
            relevantKnowledge = retrieveKnowledge(state),
        )
    }

    /** Accumulators for one assembly pass — keeps the per-kind helpers to a single parameter. */
    private class ResolvedLinks {
        val dependencies = mutableListOf<RelatedTask>()
        val documents = mutableListOf<TaskDocument>()
        val artifacts = mutableListOf<TaskArtifact>()
        val sessions = mutableListOf<RelatedSession>()
        val related = mutableListOf<RelatedLink>()
    }

    private fun addTask(link: TaskLink, out: ResolvedLinks) {
        val linked = runCatching { service.getTask(TaskId(link.targetId)) }.getOrNull()
        if (linked != null) {
            out.dependencies += RelatedTask(
                id = linked.taskId.value,
                status = linked.state.status.name,
                title = linked.state.title,
                link = link.type.name,
            )
        } else {
            out.related += RelatedLink(link.targetId, link.type.name)
        }
    }

    private suspend fun addDocument(link: TaskLink, out: ResolvedLinks) {
        val doc = documentResolver?.let { runCatching { it.resolve(link.targetId) }.getOrNull() }
        if (doc != null) {
            out.documents += TaskDocument(
                targetId = link.targetId,
                type = link.type.name,
                title = doc.title,
                path = doc.path,
                excerpt = doc.excerpt,
            )
        } else {
            out.related += RelatedLink(link.targetId, link.type.name)
        }
    }

    private suspend fun addArtifact(link: TaskLink, out: ResolvedLinks) {
        val artifact = artifactResolver?.let { runCatching { it.resolve(link.targetId) }.getOrNull() }
        if (artifact != null) {
            out.artifacts += TaskArtifact(
                targetId = link.targetId,
                type = link.type.name,
                stage = artifact.stage,
                sessionId = artifact.sessionId,
                excerpt = artifact.excerpt,
            )
        } else {
            out.related += RelatedLink(link.targetId, link.type.name)
        }
    }

    private suspend fun addSession(link: TaskLink, out: ResolvedLinks) {
        val session = sessionResolver?.let { runCatching { it.resolve(link.targetId) }.getOrNull() }
        if (session != null) {
            out.sessions += RelatedSession(
                targetId = link.targetId,
                type = link.type.name,
                status = session.status,
                intent = session.intent,
                workflowId = session.workflowId,
            )
        } else {
            out.related += RelatedLink(link.targetId, link.type.name)
        }
    }

    /**
     * The task's unfinished dependency blockers, resolved over its project board (see [TaskGraph]),
     * each labelled by the direction of the edge. Empty when the project board can't be read.
     */
    private fun computeBlockedBy(taskId: TaskId, state: TaskState): List<RelatedTask> {
        val board = state.projectId?.let { runCatching { service.list(it) }.getOrNull() } ?: return emptyList()
        return TaskGraph.unmetBlockers(Task(taskId, state), board).map { b ->
            val viaDependsOn = state.links.any {
                it.targetKind == TaskTargetKind.TASK &&
                    it.type == TaskLinkType.DEPENDS_ON &&
                    it.targetId == b.taskId.value
            }
            RelatedTask(
                id = b.taskId.value,
                status = b.state.status.name,
                title = b.state.title,
                link = if (viaDependsOn) "DEPENDS_ON" else "BLOCKS",
            )
        }
    }

    private suspend fun retrieveKnowledge(state: TaskState): List<KnowledgeHit> {
        val active = retriever ?: return emptyList()
        val query = listOfNotNull(state.title, state.goal).joinToString(" ").trim()
        if (query.isEmpty()) return emptyList()
        return runCatching { active.retrieve(query, knowledgeLimit) }.getOrDefault(emptyList())
    }

    private companion object {
        const val DEFAULT_KNOWLEDGE_LIMIT = 5
    }
}
