package com.correx.infrastructure.tools.task

import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.ProjectId
import com.correx.core.events.types.TaskLinkType
import com.correx.core.events.types.TaskNoteAuthor
import com.correx.core.events.types.TaskTargetKind
import com.correx.core.tasks.Task
import com.correx.core.tasks.TaskService
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.add
import kotlinx.serialization.json.buildJsonArray
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
import kotlinx.serialization.json.putJsonObject

/**
 * Agent-facing tool: break a large goal into a small task graph in ONE approval, instead of N
 * separate [TaskCreateTool] calls. Use when the goal has dependency seams or independent
 * review/handoff points (a thing that must land before another, or a piece worth reviewing on its
 * own); for a single coherent unit you finish in one run, use task_create. Because a session works
 * one task at a time, each node becomes a future claim — this tool decides the *shape*, claiming is
 * still lazy (task_ready → claim), never scheduled.
 *
 * Creates the tasks and their `DEPENDS_ON` edges atomically (one dedup check, one approval) and
 * returns the new ids plus which are ready to work now. An optional `parent` epic `DEPENDS_ON` every
 * child (so it completes last) and each child `IMPLEMENTS` it (provenance).
 *
 * Nested args arrive flattened (the orchestrator stringifies non-primitive tool arguments), so the
 * `tasks`/`parent` JSON is parsed here directly rather than via the list-param accessor.
 */
class TaskDecomposeTool(private val service: TaskService) : Tool, ToolExecutor {

    override val name: String = "task_decompose"
    override val description: String =
        "Break a large goal into a small task graph (parent epic + dependency-linked children) in " +
            "one shot. Use when the goal has dependency seams or independent review/handoff points; " +
            "for a single coherent unit you finish now, use task_create. Creates the tasks and their " +
            "DEPENDS_ON links in a single approval; returns the new ids and which are ready to work now."
    override val parametersSchema: JsonObject = buildJsonObject {
        put("type", "object")
        putJsonObject("properties") {
            putJsonObject("project") {
                put("type", "string")
                put("description", "Project key shared by every task, e.g. 'webui'. Ids become '<project>-<n>'.")
            }
            putJsonObject("parent") {
                put("type", "object")
                put(
                    "description",
                    "Optional umbrella/epic. It DEPENDS_ON every child (completes last); each child IMPLEMENTS it.",
                )
                putJsonObject("properties") { unitProperties() }
            }
            putJsonObject("tasks") {
                put("type", "array")
                put("description", "The units of work, in order. Declare cross-task ordering with depends_on.")
                putJsonObject("items") {
                    put("type", "object")
                    putJsonObject("properties") {
                        putJsonObject("ref") {
                            put("type", "string")
                            put("description", "Local handle other tasks cite in depends_on (e.g. 'scaffold').")
                        }
                        unitProperties()
                        putJsonObject("depends_on") {
                            put("type", "array")
                            putJsonObject("items") { put("type", "string") }
                            put(
                                "description",
                                "refs (or 0-based indices) of tasks in THIS batch that must finish first.",
                            )
                        }
                    }
                    put("required", buildJsonArray { add(JsonPrimitive("title")); add(JsonPrimitive("goal")) })
                }
            }
            putJsonObject("force") {
                put("type", "boolean")
                put("description", "Create even if a title duplicates an active task. Default false.")
            }
            putJsonObject("force_reason") {
                put("type", "string")
                put("description", "Required when force=true: why duplicates are acceptable. Recorded on each task.")
            }
        }
        put("required", buildJsonArray { add(JsonPrimitive("project")); add(JsonPrimitive("tasks")) })
    }
    override val tier: Tier = Tier.T2
    override val requiredCapabilities: Set<ToolCapability> = emptySet()

    override fun validateRequest(request: ToolRequest): ValidationResult {
        request.stringParam("project") ?: return ValidationResult.Invalid("Missing 'project' (string).")
        if (elementOf(request, "tasks") !is JsonArray) {
            return ValidationResult.Invalid("Missing 'tasks' (JSON array).")
        }
        return ValidationResult.Valid
    }

    override suspend fun execute(request: ToolRequest): ToolResult {
        val project = request.stringParam("project") ?: return fail(request, "Missing 'project' (string).")
        val specs = parseSpecs(elementOf(request, "tasks") as? JsonArray)
            ?: return fail(request, "'tasks' must be a JSON array of {title, goal, ...}.")
        if (specs.isEmpty()) return fail(request, "'tasks' must contain at least one task.")
        specs.withIndex().firstOrNull { it.value.title.isBlank() || it.value.goal.isBlank() }?.let {
            return fail(request, "task #${it.index} is missing 'title' or 'goal'.")
        }
        val parent = parseSpec(elementOf(request, "parent") as? JsonObject)
        if (parent != null && (parent.title.isBlank() || parent.goal.isBlank())) {
            return fail(request, "'parent' needs both 'title' and 'goal'.")
        }

        val adjacency = when (val e = buildEdges(specs)) {
            is Edges.Err -> return fail(request, e.message)
            is Edges.Ok -> e.adjacency
        }
        dedupFailure(request, project, specs, parent)?.let { return it }

        val pid = ProjectId(project)
        val parentTask = parent?.let {
            service.createTask(pid, it.title, it.goal, it.acceptanceCriteria, it.affectedPaths)
        }
        val children = specs.map { service.createTask(pid, it.title, it.goal, it.acceptanceCriteria, it.affectedPaths) }
        wireLinks(adjacency, children, parentTask)
        recordProvenance(request, listOfNotNull(parentTask) + children)

        val createdIds = (children + listOfNotNull(parentTask)).joinToString(",") { it.taskId.value }
        val readyIds = children.filter { service.blockers(it.taskId).isEmpty() }.map { it.taskId.value }
        return ToolResult.Success(
            invocationId = request.invocationId,
            output = render(children, parentTask, readyIds),
            metadata = mapOf("taskIds" to createdIds, "readyIds" to readyIds.joinToString(","), "project" to project),
        )
    }

    /** Create the `DEPENDS_ON` edges between children, and (if any) the parent's epic edges. */
    private suspend fun wireLinks(adjacency: List<List<Int>>, children: List<Task>, parentTask: Task?) {
        adjacency.forEachIndexed { i, deps ->
            deps.forEach { j ->
                service.link(children[i].taskId, children[j].taskId.value, TaskLinkType.DEPENDS_ON, TaskTargetKind.TASK)
            }
        }
        if (parentTask != null) {
            children.forEach { child ->
                // Parent waits on every child; child records the epic it implements.
                service.link(parentTask.taskId, child.taskId.value, TaskLinkType.DEPENDS_ON, TaskTargetKind.TASK)
                service.link(child.taskId, parentTask.taskId.value, TaskLinkType.IMPLEMENTS, TaskTargetKind.TASK)
            }
        }
    }

    private suspend fun recordProvenance(request: ToolRequest, created: List<Task>) {
        created.forEach { service.linkOriginSession(it.taskId, request) }
        if (!request.boolParam("force")) return
        val reason = request.stringParam("force_reason") ?: return
        created.forEach {
            service.addNote(
                it.taskId,
                TaskNoteAuthor.AGENT,
                "[force] created via decompose despite the duplicate guard: $reason",
            )
        }
    }

    /** Intra-batch + against-board duplicate guard, mirroring [TaskCreateTool]; force needs a reason. */
    private fun dedupFailure(
        request: ToolRequest,
        project: String,
        specs: List<TaskSpec>,
        parent: TaskSpec?,
    ): ToolResult.Failure? {
        val titles = (listOfNotNull(parent?.title) + specs.map { it.title })
        val within = titles.groupBy { normalize(it) }.filter { it.value.size > 1 }.keys
        if (within.isNotEmpty()) {
            return fail(request, "Duplicate titles within this batch: ${within.joinToString(", ")}.")
        }
        if (request.boolParam("force")) {
            return if (request.stringParam("force_reason") == null) {
                fail(request, "force=true requires 'force_reason' explaining why a duplicate is acceptable.")
            } else {
                null
            }
        }
        val pid = ProjectId(project)
        val clashes = titles.flatMap { t -> service.findDuplicates(pid, t) }.distinctBy { it.taskId.value }
        return clashes.takeIf { it.isNotEmpty() }?.let {
            val listed = it.joinToString(", ") { d ->
                "${d.taskId.value} '${d.state.title.orEmpty()}' [${d.state.status.name}]"
            }
            fail(
                request,
                "Possible duplicate(s): $listed. Reuse one (task_context/task_update), " +
                    "or force=true with force_reason.",
            )
        }
    }

    private fun render(children: List<Task>, parentTask: Task?, readyIds: List<String>): String {
        val lines = children.map { c ->
            val unmet = service.blockers(c.taskId).map { it.taskId.value }
            val state = if (unmet.isEmpty()) "READY" else "blocked by ${unmet.joinToString(", ")}"
            "- ${c.taskId.value} '${c.state.title.orEmpty()}' [$state]"
        } + listOfNotNull(
            parentTask?.let { p ->
                val unmet = service.blockers(p.taskId).map { it.taskId.value }
                "- ${p.taskId.value} '${p.state.title.orEmpty()}' (epic) [blocked by ${unmet.joinToString(", ")}]"
            },
        )
        val readyLine = if (readyIds.isEmpty()) {
            "Nothing is ready yet — check the dependency graph."
        } else {
            "Ready to work now: ${readyIds.joinToString(", ")}. Claim one to start; " +
                "the rest unblock as their dependencies complete."
        }
        return "Decomposed into ${children.size} task(s)${if (parentTask != null) " under an epic" else ""}:\n" +
            lines.joinToString("\n") + "\n" + readyLine
    }

    // --- parsing (flattened JSON args) ---

    private data class TaskSpec(
        val ref: String?,
        val title: String,
        val goal: String,
        val acceptanceCriteria: List<String>,
        val affectedPaths: List<String>,
        val deps: List<String>,
    )

    private fun parseSpecs(element: JsonArray?): List<TaskSpec>? =
        element?.map { parseSpec(it as? JsonObject) ?: return null }

    private fun parseSpec(obj: JsonObject?): TaskSpec? {
        if (obj == null) return null
        return TaskSpec(
            ref = str(obj, "ref"),
            title = str(obj, "title").orEmpty(),
            goal = str(obj, "goal").orEmpty(),
            acceptanceCriteria = strList(obj, "acceptance_criteria"),
            affectedPaths = strList(obj, "affected_paths"),
            deps = strList(obj, "depends_on"),
        )
    }

    private fun str(obj: JsonObject, key: String): String? =
        (obj[key] as? JsonPrimitive)?.content?.takeIf { it.isNotBlank() && it != "null" }

    private fun strList(obj: JsonObject, key: String): List<String> =
        (obj[key] as? JsonArray)
            ?.mapNotNull { (it as? JsonPrimitive)?.content?.takeIf(String::isNotBlank) }
            ?: emptyList()

    /** Top-level args are flattened to strings by the orchestrator, so re-parse the value as JSON. */
    private fun elementOf(request: ToolRequest, name: String) =
        request.parameters[name]?.let { runCatching { Json.parseToJsonElement(it.toString()) }.getOrNull() }

    private fun normalize(title: String): String = title.trim().lowercase().replace(Regex("\\s+"), " ")

    // --- dependency edges + cycle check ---

    private sealed interface Edges {
        data class Ok(val adjacency: List<List<Int>>) : Edges
        data class Err(val message: String) : Edges
    }

    private fun buildEdges(specs: List<TaskSpec>): Edges {
        val refToIndex = HashMap<String, Int>()
        specs.forEachIndexed { i, s -> s.ref?.let { refToIndex[it] = i } }
        val adjacency = ArrayList<List<Int>>()
        specs.forEachIndexed { i, s ->
            val deps = ArrayList<Int>()
            for (d in s.deps) {
                val j = refToIndex[d] ?: d.toIntOrNull()?.takeIf { it in specs.indices }
                    ?: return Edges.Err(
                        "task #$i ('${s.title}') depends_on '$d', which is not a ref or index in this batch.",
                    )
                if (j == i) return Edges.Err("task #$i ('${s.title}') depends on itself.")
                deps.add(j)
            }
            adjacency.add(deps.distinct())
        }
        cyclePath(adjacency)?.let { path ->
            return Edges.Err("Dependency cycle: ${path.joinToString(" -> ") { specs[it].title }}.")
        }
        return Edges.Ok(adjacency)
    }

    /** Returns a back-edge cycle as a list of task indices, or null when the graph is acyclic. */
    private fun cyclePath(adjacency: List<List<Int>>): List<Int>? {
        val color = IntArray(adjacency.size) // 0=unseen, 1=on-stack, 2=done
        val stack = ArrayList<Int>()
        fun dfs(u: Int): List<Int>? {
            color[u] = 1
            stack.add(u)
            for (v in adjacency[u]) {
                if (color[v] == 1) return stack.subList(stack.indexOf(v), stack.size).toList() + v
                if (color[v] == 0) dfs(v)?.let { return it }
            }
            color[u] = 2
            stack.removeAt(stack.lastIndex)
            return null
        }
        for (i in adjacency.indices) if (color[i] == 0) dfs(i)?.let { return it }
        return null
    }

    private fun fail(request: ToolRequest, reason: String): ToolResult.Failure =
        ToolResult.Failure(request.invocationId, reason, recoverable = false)
}

/** The fields a task spec and the parent epic share, declared once for the schema. */
private fun kotlinx.serialization.json.JsonObjectBuilder.unitProperties() {
    putJsonObject("title") { put("type", "string"); put("description", "Short imperative title.") }
    putJsonObject("goal") { put("type", "string"); put("description", "What 'done' looks like.") }
    putJsonObject("acceptance_criteria") {
        put("type", "array")
        putJsonObject("items") { put("type", "string") }
        put("description", "Concrete, verifiable criteria.")
    }
    putJsonObject("affected_paths") {
        put("type", "array")
        putJsonObject("items") { put("type", "string") }
        put("description", "Globs/paths this task is expected to touch.")
    }
}
