package com.correx.core.toolintent

import com.correx.core.events.events.ToolRequest
import com.correx.core.events.risk.RiskAction
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.toolintent.rules.PathContainmentRule
import com.correx.core.tools.contract.ParamRole
import com.correx.core.tools.contract.ToolCapability
import java.nio.file.Path
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertFalse
import kotlin.test.assertTrue

class PathContainmentRuleTest {

    private val workspace = Path.of("/work/project")
    private val rule = PathContainmentRule()

    /** Fake probe: resolveReal returns the normalized absolute path, except for the
     *  symlink test where a specific path is redirected to a privileged target. */
    private class FakeProbe(val redirect: Map<Path, Path> = emptyMap()) : WorldProbe {
        override fun exists(path: Path) = true
        override fun resolveReal(path: Path): Path =
            redirect[path.toAbsolutePath().normalize()] ?: path.toAbsolutePath().normalize()
    }

    private fun input(
        pathArg: String,
        probe: WorldProbe,
        privileged: List<Path> = listOf(Path.of("/etc"), Path.of("/usr")),
    ) = ToolCallAssessmentInput(
        request = ToolRequest(
            ToolInvocationId("i"), SessionId("s"), StageId("st"), "file_write",
            mapOf("path" to pathArg, "mode" to "0644"),
        ),
        capabilities = setOf(ToolCapability.FILE_WRITE),
        workspace = WorkspacePolicy(workspace, privileged),
        probe = probe,
    )

    @Test
    fun `applies only to file capabilities`() {
        assertTrue(rule.appliesTo(setOf(ToolCapability.FILE_READ)))
        assertTrue(rule.appliesTo(setOf(ToolCapability.FILE_WRITE)))
        assertFalse(rule.appliesTo(setOf(ToolCapability.NETWORK_ACCESS)))
        assertFalse(rule.appliesTo(emptySet()))
    }

    @Test
    fun `in-workspace path proceeds with no issues`() {
        val r = rule.assess(input("/work/project/src/A.kt", FakeProbe()))
        assertEquals(RiskAction.PROCEED, r.disposition)
        assertTrue(r.issues.isEmpty())
        assertEquals("true", r.observations.single().facts["inWorkspace"])
    }

    @Test
    fun `outside-workspace path prompts the user`() {
        val r = rule.assess(input("/tmp/scratch.txt", FakeProbe()))
        assertEquals(RiskAction.PROMPT_USER, r.disposition)
        assertEquals("PATH_OUTSIDE_WORKSPACE", r.issues.single().code)
    }

    @Test
    fun `privileged location is blocked`() {
        val r = rule.assess(input("/etc/passwd", FakeProbe()))
        assertEquals(RiskAction.BLOCK, r.disposition)
        assertEquals("PRIVILEGED_LOCATION", r.issues.single().code)
    }

    @Test
    fun `symlink inside workspace that resolves to privileged is blocked`() {
        val link = Path.of("/work/project/evil")
        val probe = FakeProbe(redirect = mapOf(link to Path.of("/etc/shadow")))
        val r = rule.assess(input("/work/project/evil", probe))
        assertEquals(RiskAction.BLOCK, r.disposition)
        assertEquals("true", r.observations.single().facts["privileged"])
    }

    @Test
    fun `non-path string parameters are ignored`() {
        val r = rule.assess(input("0644", FakeProbe())) // not path-like; "mode" also ignored
        assertEquals(RiskAction.PROCEED, r.disposition)
        assertTrue(r.observations.isEmpty())
    }

    @Test
    fun `declared PATH role flags outside-workspace path arg`() {
        val r = rule.assess(input("/tmp/scratch.txt", FakeProbe()).copy(paramRoles = mapOf("path" to ParamRole.PATH)))
        assertEquals(RiskAction.PROMPT_USER, r.disposition)
        assertEquals("PATH_OUTSIDE_WORKSPACE", r.issues.single().code)
    }

    @Test
    fun `declared PATH role ignores a slashy non-path param`() {
        val r = rule.assess(
            ToolCallAssessmentInput(
                request = ToolRequest(
                    ToolInvocationId("i"), SessionId("s"), StageId("st"), "file_write",
                    mapOf("path" to "/work/project/A.kt", "content" to "import a/b/c"),
                ),
                capabilities = setOf(ToolCapability.FILE_WRITE),
                workspace = WorkspacePolicy(workspace, listOf(Path.of("/etc"))),
                probe = FakeProbe(),
                paramRoles = mapOf("path" to ParamRole.PATH),
            ),
        )
        assertEquals(RiskAction.PROCEED, r.disposition)
        assertTrue(r.issues.isEmpty())
        assertEquals(1, r.observations.size) // only the declared path param, not the slashy content
    }

    @Test
    fun `no declared roles still flags a slashy content param via fallback`() {
        val r = rule.assess(
            ToolCallAssessmentInput(
                request = ToolRequest(
                    ToolInvocationId("i"), SessionId("s"), StageId("st"), "file_write",
                    mapOf("content" to "/tmp/evil.txt"),
                ),
                capabilities = setOf(ToolCapability.FILE_WRITE),
                workspace = WorkspacePolicy(workspace, listOf(Path.of("/etc"))),
                probe = FakeProbe(),
            ),
        )
        assertEquals(RiskAction.PROMPT_USER, r.disposition)
        assertEquals("PATH_OUTSIDE_WORKSPACE", r.issues.single().code)
    }
}
