"""Consolidated Java demo hack code.

Pre-canned optimizer, line-profiler, and testgen results for Java demo functions.
Extracted from optimizer.py, optimizer_lp.py, and testgen.py so that large string
literals do not leak into production function bodies or LLM context.
"""

from __future__ import annotations

import asyncio
import re
import uuid
from typing import TYPE_CHECKING

from aiservice.common_utils import is_host_equals_demo, should_hack_for_demo_java
from core.shared.context_helpers import group_code
from core.shared.optimizer_schemas import OptimizeResponseItemSchema, OptimizeResponseSchema
from core.shared.testgen_models import TestGenResponseSchema

if TYPE_CHECKING:
    from core.shared.testgen_models import TestGenSchema

# ---------------------------------------------------------------------------
# Helpers extracted from optimizer.py
# ---------------------------------------------------------------------------


def _extract_demo_context(source_code: str) -> tuple[str, str, str, str]:
    """Extract package, class name, exception type, and extra imports from the demo source code.

    Returns:
        Tuple of (package_declaration, class_name, throw_statement_prefix, extra_imports)

    """
    # Extract the raw code from markdown block if present
    code_match = re.search(r"```java:[^\n]*\n(.*?)```", source_code, re.DOTALL)
    raw_code = code_match.group(1) if code_match else source_code

    # Extract package
    pkg_match = re.search(r"^\s*package\s+([\w.]+)\s*;", raw_code, re.MULTILINE)
    package_decl = f"package {pkg_match.group(1)};\n" if pkg_match else ""

    # Extract class name
    class_match = re.search(r"\bclass\s+(\w+)", raw_code)
    class_name = class_match.group(1) if class_match else "FileUtils"

    # Extract exception type from the throw statement (e.g., "throw new AerospikeException(...)")
    throw_match = re.search(r"throw\s+new\s+(\w+)\s*\(", raw_code)
    exception_type = throw_match.group(1) if throw_match else "RuntimeException"

    # Collect extra imports needed for the exception type (skip standard java/javax)
    extra_imports = ""
    if exception_type != "RuntimeException":
        import_match = re.search(rf"^\s*import\s+([\w.]*\.{re.escape(exception_type)})\s*;", raw_code, re.MULTILINE)
        if import_match:
            extra_imports = f"import {import_match.group(1)};\n"

    return package_decl, class_name, exception_type, extra_imports


def _build_demo_optimizations(
    package_decl: str, class_name: str, exception_type: str, extra_imports: str
) -> list[dict[str, str]]:
    """Build 2 demo optimization candidates using the extracted class context.

    Candidate 1 (Files.readAllBytes) is the intended winner — it benchmarks fastest.
    Candidate 2 is a plausible alternative that is functionally correct but
    benchmarks slightly slower, ensuring Files.readAllBytes wins the speedup critic.
    """
    fmt = {
        "package_decl": package_decl,
        "class_name": class_name,
        "exception_type": exception_type,
        "extra_imports": extra_imports,
    }

    return [
        # Candidate 2: FileInputStream.readAllBytes() (Java 9+)
        {
            "source_code": (
                "{package_decl}"
                "\n"
                "import java.io.File;\n"
                "import java.io.FileInputStream;\n"
                "{extra_imports}"
                "\n"
                "public final class {class_name} {{\n"
                "    public static byte[] readFile(File file) {{\n"
                "        try (FileInputStream fis = new FileInputStream(file)) {{\n"
                "            return fis.readAllBytes();\n"
                "        }}\n"
                "        catch (Throwable e) {{\n"
                '            throw new {exception_type}("Failed to read " + file.getAbsolutePath(), e);\n'
                "        }}\n"
                "    }}\n"
                "}}"
            ).format(**fmt),
            "explanation": (
                "Use FileInputStream.readAllBytes() (Java 9+) to read the entire file in one call. "
                "This eliminates the manual read loop but still uses FileInputStream internally."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
        # Candidate 1: Files.readAllBytes (THE WINNER)
        {
            "source_code": (
                "{package_decl}"
                "\n"
                "import java.io.File;\n"
                "import java.nio.file.Files;\n"
                "{extra_imports}"
                "\n"
                "public final class {class_name} {{\n"
                "    public static byte[] readFile(File file) {{\n"
                "        try {{\n"
                "            return java.nio.file.Files.readAllBytes(file.toPath());\n"
                "        }}\n"
                "        catch (Throwable e) {{\n"
                '            throw new {exception_type}("Failed to read " + file.getAbsolutePath(), e);\n'
                "        }}\n"
                "    }}\n"
                "}}"
            ).format(**fmt),
            "explanation": (
                "Replace manual FileInputStream read loop with java.nio.file.Files.readAllBytes(). "
                "This NIO method is optimized at the JDK level for direct file-to-byte-array transfer, "
                "eliminating manual buffering and loop overhead."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
    ]


def _build_host_equals_demo_optimizations(source_code: str) -> list[dict[str, str]]:
    """Build 5 optimization candidates for Host.equals by reordering comparisons.

    Candidate 1 (port-first early return) is the intended winner — comparing the
    primitive int port before the String name avoids unnecessary method dispatch.
    """
    code_match = re.search(r"```java:[^\n]*\n(.*?)```", source_code, re.DOTALL)
    raw_code = code_match.group(1) if code_match else source_code

    # Match: return this.name.equals(other.name) && this.port == other.port;
    original_stmt = re.compile(
        r"(\s*)return\s+this\.name\.equals\(other\.name\)\s*&&\s*this\.port\s*==\s*other\.port\s*;"
    )

    match = original_stmt.search(raw_code)
    if not match:
        return [
            {
                "source_code": raw_code,
                "explanation": "No optimization applicable.",
                "optimization_id": str(uuid.uuid4()),
            }
        ]

    indent = match.group(1)
    inner = indent + "    "

    def replace_with(replacement: str) -> str:
        return original_stmt.sub(replacement, raw_code)

    return [
        # Candidate 1 (WINNER): Port-first early return
        {
            "source_code": replace_with(
                f"{indent}// Compare primitive port first to avoid unnecessary string equals calls.\n"
                f"{indent}if (this.port != other.port) {{\n"
                f"{inner}return false;\n"
                f"{indent}}}\n"
                f"{indent}return this.name.equals(other.name);"
            ),
            "explanation": (
                "Compare primitive port first to avoid unnecessary string equals calls. "
                "Integer comparison is a single CPU instruction, while String.equals() "
                "involves method dispatch and potential character-by-character comparison."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
        # Candidate 2: Reordered conjunction (port first in &&)
        {
            "source_code": replace_with(f"{indent}return this.port == other.port && this.name.equals(other.name);"),
            "explanation": (
                "Reorder the conjunction to evaluate the cheaper primitive int comparison first. "
                "Short-circuit evaluation skips String.equals() when ports differ."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
        # Candidate 3: Port-first with Objects.equals for null safety
        {
            "source_code": replace_with(
                f"{indent}if (this.port != other.port) {{\n"
                f"{inner}return false;\n"
                f"{indent}}}\n"
                f"{indent}return java.util.Objects.equals(this.name, other.name);"
            ),
            "explanation": (
                "Check port first (cheap primitive comparison), then use Objects.equals() "
                "for null-safe name comparison. Adds safety at slight method-call overhead."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
        # Candidate 4: Ternary with port-first guard
        {
            "source_code": replace_with(
                f"{indent}return this.port == other.port ? this.name.equals(other.name) : false;"
            ),
            "explanation": (
                "Use a ternary to short-circuit on port mismatch. "
                "Evaluates the cheap int comparison first, only calling String.equals() when ports match."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
        # Candidate 5: Explicit null guard + port first
        {
            "source_code": replace_with(
                f"{indent}if (this.port != other.port) {{\n"
                f"{inner}return false;\n"
                f"{indent}}}\n"
                f"{indent}if (this.name == null) {{\n"
                f"{inner}return other.name == null;\n"
                f"{indent}}}\n"
                f"{indent}return this.name.equals(other.name);"
            ),
            "explanation": (
                "Guard on port first, then add explicit null handling for the name field "
                "before delegating to String.equals(). Avoids potential NullPointerException."
            ),
            "optimization_id": str(uuid.uuid4()),
        },
    ]


# ---------------------------------------------------------------------------
# Demo optimizer response  (from optimizer.py)
# ---------------------------------------------------------------------------


async def hack_for_demo_java(source_code: str) -> OptimizeResponseSchema:
    # Extract file path from markdown source (```java:path/to/File.java)
    file_path_match = re.search(r"```java:([^\n]+)", source_code)
    file_name = file_path_match.group(1).strip() if file_path_match else "Source.java"

    if is_host_equals_demo(source_code):
        optimizations = _build_host_equals_demo_optimizations(source_code)
    else:
        # Extract class context dynamically from the source code
        package_decl, class_name, exception_type, extra_imports = _extract_demo_context(source_code)
        optimizations = _build_demo_optimizations(package_decl, class_name, exception_type, extra_imports)

    response_list: list[OptimizeResponseItemSchema] = [
        OptimizeResponseItemSchema(
            explanation=opt["explanation"],
            optimization_id=opt["optimization_id"],
            source_code=group_code({file_name: opt["source_code"]}, language="java"),
        )
        for opt in optimizations
    ]
    await asyncio.sleep(5)
    return OptimizeResponseSchema(optimizations=response_list)


# ---------------------------------------------------------------------------
# Demo LP optimizer response  (from optimizer_lp.py)
# ---------------------------------------------------------------------------


async def hack_for_demo_java_lp(source_code: str) -> OptimizeResponseSchema:
    """Return pre-canned line-profiler optimization results for the Java demo function."""
    file_path_match = re.search(r"```java:([^\n]+)", source_code)
    file_name = file_path_match.group(1).strip() if file_path_match else "Source.java"

    if is_host_equals_demo(source_code):
        optimizations = _build_host_equals_demo_optimizations(source_code)
    else:
        package_decl, class_name, exception_type, extra_imports = _extract_demo_context(source_code)
        optimizations = _build_demo_optimizations(package_decl, class_name, exception_type, extra_imports)

    response_list: list[OptimizeResponseItemSchema] = [
        OptimizeResponseItemSchema(
            explanation=opt["explanation"],
            optimization_id=opt["optimization_id"],
            source_code=group_code({file_name: opt["source_code"]}, language="java"),
        )
        for opt in optimizations
    ]
    await asyncio.sleep(5)
    return OptimizeResponseSchema(optimizations=response_list)


# ---------------------------------------------------------------------------
# Demo test builders  (from testgen.py)
# ---------------------------------------------------------------------------


def _build_demo_test_source_0(package_name: str, class_name: str, test_framework: str) -> str:
    """Build demo test source 0, adapting to the target's package, class, and test framework.

    File creation is in @Before/@BeforeEach so it runs once, outside the instrumentation's
    inner loop. Test methods only contain readFile calls so every inner iteration succeeds
    and the benchmark measures pure readFile performance.
    """
    module_path = f"{package_name}.{class_name}" if package_name else class_name
    test_class_name = f"{class_name}Test"

    if test_framework == "junit4":
        return (f"package {package_name};\n" if package_name else "") + (
            "\n"
            "import org.junit.Before;\n"
            "import org.junit.Test;\n"
            "import org.junit.Rule;\n"
            "import org.junit.rules.TemporaryFolder;\n"
            "import static org.junit.Assert.*;\n"
            "\n"
            "import java.io.File;\n"
            "import java.io.FileOutputStream;\n"
            "\n"
            f"import {module_path};\n"
            "\n"
            f"public class {test_class_name} {{\n"
            "\n"
            "    @Rule\n"
            "    public TemporaryFolder tempFolder = new TemporaryFolder();\n"
            "\n"
            "    private File smallFile;\n"
            "    private byte[] expectedSmall;\n"
            "    private File mediumFile;\n"
            "    private byte[] expectedMedium;\n"
            "    private File largeFile;\n"
            "    private byte[] expectedLarge;\n"
            "\n"
            "    @Before\n"
            "    public void setUp() throws Exception {\n"
            '        smallFile = tempFolder.newFile("small.txt");\n'
            '        expectedSmall = "Hello, World!".getBytes();\n'
            "        try (FileOutputStream out = new FileOutputStream(smallFile)) {\n"
            "            out.write(expectedSmall);\n"
            "        }\n"
            "\n"
            '        mediumFile = tempFolder.newFile("medium.dat");\n'
            "        expectedMedium = new byte[256 * 1024];\n"
            "        for (int i = 0; i < expectedMedium.length; i++) {\n"
            "            expectedMedium[i] = (byte) (i % 251);\n"
            "        }\n"
            "        try (FileOutputStream out = new FileOutputStream(mediumFile)) {\n"
            "            out.write(expectedMedium);\n"
            "        }\n"
            "\n"
            '        largeFile = tempFolder.newFile("large.dat");\n'
            "        expectedLarge = new byte[1024 * 1024];\n"
            "        for (int i = 0; i < expectedLarge.length; i++) {\n"
            "            expectedLarge[i] = (byte) (i % 256);\n"
            "        }\n"
            "        try (FileOutputStream out = new FileOutputStream(largeFile)) {\n"
            "            out.write(expectedLarge);\n"
            "        }\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testReadSmallFile() throws Exception {\n"
            f"        byte[] result = {class_name}.readFile(smallFile);\n"
            "        assertArrayEquals(expectedSmall, result);\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testReadMediumFile() throws Exception {\n"
            f"        byte[] result = {class_name}.readFile(mediumFile);\n"
            "        assertArrayEquals(expectedMedium, result);\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testReadLargeFile() throws Exception {\n"
            f"        byte[] result = {class_name}.readFile(largeFile);\n"
            "        assertArrayEquals(expectedLarge, result);\n"
            "    }\n"
            "}\n"
        )
    # JUnit 5
    return (f"package {package_name};\n" if package_name else "") + (
        "\n"
        "import org.junit.jupiter.api.BeforeEach;\n"
        "import org.junit.jupiter.api.Test;\n"
        "import org.junit.jupiter.api.DisplayName;\n"
        "import org.junit.jupiter.api.io.TempDir;\n"
        "import static org.junit.jupiter.api.Assertions.*;\n"
        "\n"
        "import java.io.File;\n"
        "import java.io.FileOutputStream;\n"
        "import java.nio.file.Path;\n"
        "\n"
        f"import {module_path};\n"
        "\n"
        f"class {test_class_name} {{\n"
        "\n"
        "    @TempDir\n"
        "    Path tempDir;\n"
        "\n"
        "    private File smallFile;\n"
        "    private byte[] expectedSmall;\n"
        "    private File mediumFile;\n"
        "    private byte[] expectedMedium;\n"
        "    private File largeFile;\n"
        "    private byte[] expectedLarge;\n"
        "\n"
        "    @BeforeEach\n"
        "    void setUp() throws Exception {\n"
        '        smallFile = tempDir.resolve("small.txt").toFile();\n'
        '        expectedSmall = "Hello, World!".getBytes();\n'
        "        try (FileOutputStream out = new FileOutputStream(smallFile)) {\n"
        "            out.write(expectedSmall);\n"
        "        }\n"
        "\n"
        '        mediumFile = tempDir.resolve("medium.dat").toFile();\n'
        "        expectedMedium = new byte[256 * 1024];\n"
        "        for (int i = 0; i < expectedMedium.length; i++) {\n"
        "            expectedMedium[i] = (byte) (i % 251);\n"
        "        }\n"
        "        try (FileOutputStream out = new FileOutputStream(mediumFile)) {\n"
        "            out.write(expectedMedium);\n"
        "        }\n"
        "\n"
        '        largeFile = tempDir.resolve("large.dat").toFile();\n'
        "        expectedLarge = new byte[1024 * 1024];\n"
        "        for (int i = 0; i < expectedLarge.length; i++) {\n"
        "            expectedLarge[i] = (byte) (i % 256);\n"
        "        }\n"
        "        try (FileOutputStream out = new FileOutputStream(largeFile)) {\n"
        "            out.write(expectedLarge);\n"
        "        }\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Read a small file")\n'
        "    void testReadSmallFile() throws Exception {\n"
        f"        byte[] result = {class_name}.readFile(smallFile);\n"
        "        assertArrayEquals(expectedSmall, result);\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Read a 256KB file")\n'
        "    void testReadMediumFile() throws Exception {\n"
        f"        byte[] result = {class_name}.readFile(mediumFile);\n"
        "        assertArrayEquals(expectedMedium, result);\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Read a 1MB file")\n'
        "    void testReadLargeFile() throws Exception {\n"
        f"        byte[] result = {class_name}.readFile(largeFile);\n"
        "        assertArrayEquals(expectedLarge, result);\n"
        "    }\n"
        "}\n"
    )


def _build_demo_test_source_1(package_name: str, class_name: str, test_framework: str) -> str:
    """Build demo test source 1, adapting to the target's package, class, and test framework.

    Same @Before pattern as source_0: file creation outside the timed test body.
    Complementary file sizes to source_0.
    """
    module_path = f"{package_name}.{class_name}" if package_name else class_name
    test_class_name = f"{class_name}Test"

    if test_framework == "junit4":
        return (f"package {package_name};\n" if package_name else "") + (
            "\n"
            "import org.junit.Before;\n"
            "import org.junit.Test;\n"
            "import org.junit.Rule;\n"
            "import org.junit.rules.TemporaryFolder;\n"
            "import static org.junit.Assert.*;\n"
            "\n"
            "import java.io.File;\n"
            "import java.io.FileOutputStream;\n"
            "import java.util.Arrays;\n"
            "\n"
            f"import {module_path};\n"
            "\n"
            f"public class {test_class_name} {{\n"
            "\n"
            "    @Rule\n"
            "    public TemporaryFolder tempFolder = new TemporaryFolder();\n"
            "\n"
            "    private File patternFile;\n"
            "    private byte[] expectedPattern;\n"
            "    private File halfMegFile;\n"
            "    private byte[] expectedHalfMeg;\n"
            "    private File twoMegFile;\n"
            "    private byte[] expectedTwoMeg;\n"
            "\n"
            "    @Before\n"
            "    public void setUp() throws Exception {\n"
            '        patternFile = tempFolder.newFile("pattern.dat");\n'
            "        expectedPattern = new byte[128 * 1024];\n"
            "        for (int i = 0; i < expectedPattern.length; i++) {\n"
            "            expectedPattern[i] = (byte) (i % 7);\n"
            "        }\n"
            "        try (FileOutputStream out = new FileOutputStream(patternFile)) {\n"
            "            out.write(expectedPattern);\n"
            "        }\n"
            "\n"
            '        halfMegFile = tempFolder.newFile("half_meg.dat");\n'
            "        expectedHalfMeg = new byte[512 * 1024];\n"
            "        Arrays.fill(expectedHalfMeg, (byte) 0xCD);\n"
            "        try (FileOutputStream out = new FileOutputStream(halfMegFile)) {\n"
            "            out.write(expectedHalfMeg);\n"
            "        }\n"
            "\n"
            '        twoMegFile = tempFolder.newFile("two_meg.dat");\n'
            "        expectedTwoMeg = new byte[2 * 1024 * 1024];\n"
            "        for (int i = 0; i < expectedTwoMeg.length; i++) {\n"
            "            expectedTwoMeg[i] = (byte) (i % 199);\n"
            "        }\n"
            "        try (FileOutputStream out = new FileOutputStream(twoMegFile)) {\n"
            "            out.write(expectedTwoMeg);\n"
            "        }\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testReadBinaryPattern() throws Exception {\n"
            f"        byte[] result = {class_name}.readFile(patternFile);\n"
            "        assertArrayEquals(expectedPattern, result);\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testReadHalfMegFile() throws Exception {\n"
            f"        byte[] result = {class_name}.readFile(halfMegFile);\n"
            "        assertArrayEquals(expectedHalfMeg, result);\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testReadTwoMegFile() throws Exception {\n"
            f"        byte[] result = {class_name}.readFile(twoMegFile);\n"
            "        assertArrayEquals(expectedTwoMeg, result);\n"
            "    }\n"
            "}\n"
        )
    # JUnit 5
    return (f"package {package_name};\n" if package_name else "") + (
        "\n"
        "import org.junit.jupiter.api.BeforeEach;\n"
        "import org.junit.jupiter.api.Test;\n"
        "import org.junit.jupiter.api.DisplayName;\n"
        "import org.junit.jupiter.api.io.TempDir;\n"
        "import static org.junit.jupiter.api.Assertions.*;\n"
        "\n"
        "import java.io.File;\n"
        "import java.io.FileOutputStream;\n"
        "import java.nio.file.Path;\n"
        "import java.util.Arrays;\n"
        "\n"
        f"import {module_path};\n"
        "\n"
        f"class {test_class_name} {{\n"
        "\n"
        "    @TempDir\n"
        "    Path tempDir;\n"
        "\n"
        "    private File patternFile;\n"
        "    private byte[] expectedPattern;\n"
        "    private File halfMegFile;\n"
        "    private byte[] expectedHalfMeg;\n"
        "    private File twoMegFile;\n"
        "    private byte[] expectedTwoMeg;\n"
        "\n"
        "    @BeforeEach\n"
        "    void setUp() throws Exception {\n"
        '        patternFile = tempDir.resolve("pattern.dat").toFile();\n'
        "        expectedPattern = new byte[128 * 1024];\n"
        "        for (int i = 0; i < expectedPattern.length; i++) {\n"
        "            expectedPattern[i] = (byte) (i % 7);\n"
        "        }\n"
        "        try (FileOutputStream out = new FileOutputStream(patternFile)) {\n"
        "            out.write(expectedPattern);\n"
        "        }\n"
        "\n"
        '        halfMegFile = tempDir.resolve("half_meg.dat").toFile();\n'
        "        expectedHalfMeg = new byte[512 * 1024];\n"
        "        Arrays.fill(expectedHalfMeg, (byte) 0xCD);\n"
        "        try (FileOutputStream out = new FileOutputStream(halfMegFile)) {\n"
        "            out.write(expectedHalfMeg);\n"
        "        }\n"
        "\n"
        '        twoMegFile = tempDir.resolve("two_meg.dat").toFile();\n'
        "        expectedTwoMeg = new byte[2 * 1024 * 1024];\n"
        "        for (int i = 0; i < expectedTwoMeg.length; i++) {\n"
        "            expectedTwoMeg[i] = (byte) (i % 199);\n"
        "        }\n"
        "        try (FileOutputStream out = new FileOutputStream(twoMegFile)) {\n"
        "            out.write(expectedTwoMeg);\n"
        "        }\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Read 128KB binary pattern file")\n'
        "    void testReadBinaryPattern() throws Exception {\n"
        f"        byte[] result = {class_name}.readFile(patternFile);\n"
        "        assertArrayEquals(expectedPattern, result);\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Read 512KB file")\n'
        "    void testReadHalfMegFile() throws Exception {\n"
        f"        byte[] result = {class_name}.readFile(halfMegFile);\n"
        "        assertArrayEquals(expectedHalfMeg, result);\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Read 2MB file")\n'
        "    void testReadTwoMegFile() throws Exception {\n"
        f"        byte[] result = {class_name}.readFile(twoMegFile);\n"
        "        assertArrayEquals(expectedTwoMeg, result);\n"
        "    }\n"
        "}\n"
    )


def _build_host_equals_demo_test_source_0(package_name: str, class_name: str, test_framework: str) -> str:
    """Build demo test source 0 for Host.equals, adapting to the target's package, class, and test framework."""
    module_path = f"{package_name}.{class_name}" if package_name else class_name
    test_class_name = f"{class_name}Test"

    if test_framework == "junit4":
        return (f"package {package_name};\n" if package_name else "") + (
            "\n"
            "import org.junit.Before;\n"
            "import org.junit.Test;\n"
            "import static org.junit.Assert.*;\n"
            "\n"
            f"import {module_path};\n"
            "\n"
            "/**\n"
            f" * Unit tests for {module_path}.equals(...)\n"
            " */\n"
            f"public class {test_class_name} {{\n"
            f"    private {class_name} defaultHost;\n"
            "\n"
            "    @Before\n"
            "    public void setUp() {\n"
            f'        defaultHost = new {class_name}("localhost", 3000);\n'
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_SameInstance_ReturnsTrue() {\n"
            "        assertTrue(defaultHost.equals(defaultHost));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_EqualNameAndPort_IgnoringTls_ReturnsTrue() {\n"
            f'        {class_name} withTls = new {class_name}("localhost", "server-cert", 3000);\n'
            "        // Both directions should be true (symmetry)\n"
            "        assertTrue(defaultHost.equals(withTls) && withTls.equals(defaultHost));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_DifferentPort_ReturnsFalse() {\n"
            f'        {class_name} other = new {class_name}("localhost", 3001);\n'
            "        assertFalse(defaultHost.equals(other));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_DifferentName_ReturnsFalse() {\n"
            f'        {class_name} other = new {class_name}("otherhost", 3000);\n'
            "        assertFalse(defaultHost.equals(other));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_NullArgument_ReturnsFalse() {\n"
            "        assertFalse(defaultHost.equals(null));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_DifferentClass_ReturnsFalse() {\n"
            '        Object notAHost = "I am not a Host";\n'
            "        assertFalse(defaultHost.equals(notAHost));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_EmptyNameBoth_ReturnsTrue() {\n"
            f'        {class_name} a = new {class_name}("", 0);\n'
            f'        {class_name} b = new {class_name}("", null, 0);\n'
            "        assertTrue(a.equals(b));\n"
            "    }\n"
            "\n"
            "    @Test(expected = NullPointerException.class)\n"
            "    public void testEquals_ThisNameNull_ThrowsNullPointerException() {\n"
            "        // When this.name is null, equals calls this.name.equals(...), which throws NPE.\n"
            f"        {class_name} thisHasNullName = new {class_name}(null, 100);\n"
            f'        {class_name} other = new {class_name}("something", 100);\n'
            "        thisHasNullName.equals(other);\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_OtherNameNull_ReturnsFalse() {\n"
            f"        {class_name} otherHasNullName = new {class_name}(null, 200);\n"
            f'        {class_name} normal = new {class_name}("name", 200);\n'
            '        // "name".equals(null) returns false; no exception expected.\n'
            "        assertFalse(normal.equals(otherHasNullName));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_MaxIntPort_ReturnsTrue() {\n"
            f'        {class_name} a = new {class_name}("host", Integer.MAX_VALUE);\n'
            f'        {class_name} b = new {class_name}("host", Integer.MAX_VALUE);\n'
            "        assertTrue(a.equals(b));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_MinIntPort_ReturnsTrue() {\n"
            f'        {class_name} a = new {class_name}("host", Integer.MIN_VALUE);\n'
            f'        {class_name} b = new {class_name}("host", Integer.MIN_VALUE);\n'
            "        assertTrue(a.equals(b));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_MaxAndDifferentPort_ReturnsFalse() {\n"
            f'        {class_name} a = new {class_name}("host", Integer.MAX_VALUE);\n'
            f'        {class_name} b = new {class_name}("host", Integer.MAX_VALUE - 1);\n'
            "        assertFalse(a.equals(b));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEquals_MultipleEqualInstances_ReturnsTrue() {\n"
            f'        {class_name} a = new {class_name}("perf-host", 4000);\n'
            f'        {class_name} b = new {class_name}("perf-host", 4000);\n'
            f'        {class_name} c = new {class_name}("perf-host", "tls-name", 4000);\n'
            "        assertTrue(a.equals(b));\n"
            "        assertTrue(b.equals(c));\n"
            "        assertTrue(a.equals(c));\n"
            "    }\n"
            "}\n"
        )
    # JUnit 5
    return (f"package {package_name};\n" if package_name else "") + (
        "\n"
        "import org.junit.jupiter.api.BeforeEach;\n"
        "import org.junit.jupiter.api.Test;\n"
        "import org.junit.jupiter.api.DisplayName;\n"
        "import static org.junit.jupiter.api.Assertions.*;\n"
        "\n"
        f"import {module_path};\n"
        "\n"
        "/**\n"
        f" * Unit tests for {module_path}.equals(...)\n"
        " */\n"
        f"class {test_class_name} {{\n"
        f"    private {class_name} defaultHost;\n"
        "\n"
        "    @BeforeEach\n"
        "    void setUp() {\n"
        f'        defaultHost = new {class_name}("localhost", 3000);\n'
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Same instance returns true")\n'
        "    void testEquals_SameInstance_ReturnsTrue() {\n"
        "        assertTrue(defaultHost.equals(defaultHost));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Equal name and port ignoring TLS returns true")\n'
        "    void testEquals_EqualNameAndPort_IgnoringTls_ReturnsTrue() {\n"
        f'        {class_name} withTls = new {class_name}("localhost", "server-cert", 3000);\n'
        "        assertTrue(defaultHost.equals(withTls) && withTls.equals(defaultHost));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different port returns false")\n'
        "    void testEquals_DifferentPort_ReturnsFalse() {\n"
        f'        {class_name} other = new {class_name}("localhost", 3001);\n'
        "        assertFalse(defaultHost.equals(other));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different name returns false")\n'
        "    void testEquals_DifferentName_ReturnsFalse() {\n"
        f'        {class_name} other = new {class_name}("otherhost", 3000);\n'
        "        assertFalse(defaultHost.equals(other));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Null argument returns false")\n'
        "    void testEquals_NullArgument_ReturnsFalse() {\n"
        "        assertFalse(defaultHost.equals(null));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different class returns false")\n'
        "    void testEquals_DifferentClass_ReturnsFalse() {\n"
        '        Object notAHost = "I am not a Host";\n'
        "        assertFalse(defaultHost.equals(notAHost));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Empty name both returns true")\n'
        "    void testEquals_EmptyNameBoth_ReturnsTrue() {\n"
        f'        {class_name} a = new {class_name}("", 0);\n'
        f'        {class_name} b = new {class_name}("", null, 0);\n'
        "        assertTrue(a.equals(b));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Null this.name throws NullPointerException")\n'
        "    void testEquals_ThisNameNull_ThrowsNullPointerException() {\n"
        f"        {class_name} thisHasNullName = new {class_name}(null, 100);\n"
        f'        {class_name} other = new {class_name}("something", 100);\n'
        "        assertThrows(NullPointerException.class, () -> thisHasNullName.equals(other));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Null other.name returns false")\n'
        "    void testEquals_OtherNameNull_ReturnsFalse() {\n"
        f"        {class_name} otherHasNullName = new {class_name}(null, 200);\n"
        f'        {class_name} normal = new {class_name}("name", 200);\n'
        "        assertFalse(normal.equals(otherHasNullName));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Max int port returns true")\n'
        "    void testEquals_MaxIntPort_ReturnsTrue() {\n"
        f'        {class_name} a = new {class_name}("host", Integer.MAX_VALUE);\n'
        f'        {class_name} b = new {class_name}("host", Integer.MAX_VALUE);\n'
        "        assertTrue(a.equals(b));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Min int port returns true")\n'
        "    void testEquals_MinIntPort_ReturnsTrue() {\n"
        f'        {class_name} a = new {class_name}("host", Integer.MIN_VALUE);\n'
        f'        {class_name} b = new {class_name}("host", Integer.MIN_VALUE);\n'
        "        assertTrue(a.equals(b));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Max vs different port returns false")\n'
        "    void testEquals_MaxAndDifferentPort_ReturnsFalse() {\n"
        f'        {class_name} a = new {class_name}("host", Integer.MAX_VALUE);\n'
        f'        {class_name} b = new {class_name}("host", Integer.MAX_VALUE - 1);\n'
        "        assertFalse(a.equals(b));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Multiple equal instances with transitivity")\n'
        "    void testEquals_MultipleEqualInstances_ReturnsTrue() {\n"
        f'        {class_name} a = new {class_name}("perf-host", 4000);\n'
        f'        {class_name} b = new {class_name}("perf-host", 4000);\n'
        f'        {class_name} c = new {class_name}("perf-host", "tls-name", 4000);\n'
        "        assertTrue(a.equals(b));\n"
        "        assertTrue(b.equals(c));\n"
        "        assertTrue(a.equals(c));\n"
        "    }\n"
        "}\n"
    )


def _build_host_equals_demo_test_source_1(package_name: str, class_name: str, test_framework: str) -> str:
    """Build demo test source 1 for Host.equals, adapting to the target's package, class, and test framework."""
    module_path = f"{package_name}.{class_name}" if package_name else class_name
    test_class_name = f"{class_name}Test"

    if test_framework == "junit4":
        return (f"package {package_name};\n" if package_name else "") + (
            "\n"
            "import org.junit.Before;\n"
            "import org.junit.Test;\n"
            "import static org.junit.Assert.*;\n"
            "\n"
            f"import {module_path};\n"
            "\n"
            f"public class {test_class_name} {{\n"
            f"    private {class_name} hostSimple;\n"
            f"    private {class_name} hostWithTls;\n"
            "\n"
            "    @Before\n"
            "    public void setUp() {\n"
            f'        hostSimple = new {class_name}("server.example.com", 3000);\n'
            f'        hostWithTls = new {class_name}("server.example.com", "tls.server.example.com", 3000);\n'
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testSameReference_True() {\n"
            "        // same instance should be equal to itself\n"
            "        assertTrue(hostSimple.equals(hostSimple));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEqualNameAndPort_True() {\n"
            "        // two distinct instances with same name and port (tls ignored) are equal\n"
            f'        {class_name} a = new {class_name}("db1", 4000);\n'
            f'        {class_name} b = new {class_name}("db1", "tlsNameDifferent", 4000);\n'
            "        assertTrue(a.equals(b));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testDifferentTlsIgnored_True() {\n"
            "        // tlsName is ignored for equality\n"
            "        assertTrue(hostSimple.equals(hostWithTls));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testDifferentName_False() {\n"
            f'        {class_name} otherName = new {class_name}("other.example.com", 3000);\n'
            "        assertFalse(hostSimple.equals(otherName));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testDifferentPort_False() {\n"
            f'        {class_name} otherPort = new {class_name}("server.example.com", 3001);\n'
            "        assertFalse(hostSimple.equals(otherPort));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testNullComparison_False() {\n"
            "        // equals should return false when compared to null\n"
            "        assertFalse(hostSimple.equals(null));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testDifferentClass_False() {\n"
            "        // equals should return false when compared to an object of another class\n"
            '        Object notAHost = "server.example.com:3000";\n'
            "        assertFalse(hostSimple.equals(notAHost));\n"
            "    }\n"
            "\n"
            "    @Test(expected = NullPointerException.class)\n"
            "    public void testNameNull_ThrowsNullPointerException() {\n"
            "        // If this.name is null, equals tries to call this.name.equals(...) and will NPE.\n"
            f"        {class_name} nullNameHost = new {class_name}(null, 3000);\n"
            f"        {class_name} otherNullNameHost = new {class_name}(null, 3000);\n"
            "        // This invocation should throw NPE because this.name is null\n"
            "        nullNameHost.equals(otherNullNameHost);\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testOtherNameNull_False() {\n"
            "        // If other.name is null but this.name is non-null, equals should return false\n"
            f"        {class_name} otherNullName = new {class_name}(null, 3000);\n"
            "        assertFalse(hostSimple.equals(otherNullName));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testPortBoundary_ZeroAndMax_True() {\n"
            f'        {class_name} lowA = new {class_name}("edge", 0);\n'
            f'        {class_name} lowB = new {class_name}("edge", 0);\n'
            "        assertTrue(lowA.equals(lowB));\n"
            "\n"
            f'        {class_name} highA = new {class_name}("edge", 65535);\n'
            f'        {class_name} highB = new {class_name}("edge", 65535);\n'
            "        assertTrue(highA.equals(highB));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testPortBoundary_DifferentPorts_False() {\n"
            f'        {class_name} low = new {class_name}("edge", 0);\n'
            f'        {class_name} high = new {class_name}("edge", 65535);\n'
            "        assertFalse(low.equals(high));\n"
            "    }\n"
            "\n"
            "    @Test\n"
            "    public void testEqualityWithVariousNames() {\n"
            f'        {class_name} a1 = new {class_name}("alpha", 1000);\n'
            f'        {class_name} b1 = new {class_name}("alpha", "tls-alpha", 1000);\n'
            "        assertTrue(a1.equals(b1));\n"
            "\n"
            f'        {class_name} a2 = new {class_name}("beta", 2000);\n'
            f'        {class_name} b2 = new {class_name}("beta", "tls-beta", 2000);\n'
            "        assertTrue(a2.equals(b2));\n"
            "\n"
            f'        {class_name} a3 = new {class_name}("gamma", 3000);\n'
            f'        {class_name} b3 = new {class_name}("delta", 3000);\n'
            "        assertFalse(a3.equals(b3));\n"
            "    }\n"
            "}\n"
        )
    # JUnit 5
    return (f"package {package_name};\n" if package_name else "") + (
        "\n"
        "import org.junit.jupiter.api.BeforeEach;\n"
        "import org.junit.jupiter.api.Test;\n"
        "import org.junit.jupiter.api.DisplayName;\n"
        "import static org.junit.jupiter.api.Assertions.*;\n"
        "\n"
        f"import {module_path};\n"
        "\n"
        f"class {test_class_name} {{\n"
        f"    private {class_name} hostSimple;\n"
        f"    private {class_name} hostWithTls;\n"
        "\n"
        "    @BeforeEach\n"
        "    void setUp() {\n"
        f'        hostSimple = new {class_name}("server.example.com", 3000);\n'
        f'        hostWithTls = new {class_name}("server.example.com", "tls.server.example.com", 3000);\n'
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Same reference returns true")\n'
        "    void testSameReference_True() {\n"
        "        assertTrue(hostSimple.equals(hostSimple));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Equal name and port with different TLS returns true")\n'
        "    void testEqualNameAndPort_True() {\n"
        f'        {class_name} a = new {class_name}("db1", 4000);\n'
        f'        {class_name} b = new {class_name}("db1", "tlsNameDifferent", 4000);\n'
        "        assertTrue(a.equals(b));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different TLS name is ignored")\n'
        "    void testDifferentTlsIgnored_True() {\n"
        "        assertTrue(hostSimple.equals(hostWithTls));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different name returns false")\n'
        "    void testDifferentName_False() {\n"
        f'        {class_name} otherName = new {class_name}("other.example.com", 3000);\n'
        "        assertFalse(hostSimple.equals(otherName));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different port returns false")\n'
        "    void testDifferentPort_False() {\n"
        f'        {class_name} otherPort = new {class_name}("server.example.com", 3001);\n'
        "        assertFalse(hostSimple.equals(otherPort));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Null comparison returns false")\n'
        "    void testNullComparison_False() {\n"
        "        assertFalse(hostSimple.equals(null));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different class returns false")\n'
        "    void testDifferentClass_False() {\n"
        '        Object notAHost = "server.example.com:3000";\n'
        "        assertFalse(hostSimple.equals(notAHost));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Null name throws NullPointerException")\n'
        "    void testNameNull_ThrowsNullPointerException() {\n"
        f"        {class_name} nullNameHost = new {class_name}(null, 3000);\n"
        f"        {class_name} otherNullNameHost = new {class_name}(null, 3000);\n"
        "        assertThrows(NullPointerException.class, () -> nullNameHost.equals(otherNullNameHost));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Other name null returns false")\n'
        "    void testOtherNameNull_False() {\n"
        f"        {class_name} otherNullName = new {class_name}(null, 3000);\n"
        "        assertFalse(hostSimple.equals(otherNullName));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Port boundary zero and max")\n'
        "    void testPortBoundary_ZeroAndMax_True() {\n"
        f'        {class_name} lowA = new {class_name}("edge", 0);\n'
        f'        {class_name} lowB = new {class_name}("edge", 0);\n'
        "        assertTrue(lowA.equals(lowB));\n"
        "\n"
        f'        {class_name} highA = new {class_name}("edge", 65535);\n'
        f'        {class_name} highB = new {class_name}("edge", 65535);\n'
        "        assertTrue(highA.equals(highB));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Different boundary ports returns false")\n'
        "    void testPortBoundary_DifferentPorts_False() {\n"
        f'        {class_name} low = new {class_name}("edge", 0);\n'
        f'        {class_name} high = new {class_name}("edge", 65535);\n'
        "        assertFalse(low.equals(high));\n"
        "    }\n"
        "\n"
        "    @Test\n"
        '    @DisplayName("Equality with various host names")\n'
        "    void testEqualityWithVariousNames() {\n"
        f'        {class_name} a1 = new {class_name}("alpha", 1000);\n'
        f'        {class_name} b1 = new {class_name}("alpha", "tls-alpha", 1000);\n'
        "        assertTrue(a1.equals(b1));\n"
        "\n"
        f'        {class_name} a2 = new {class_name}("beta", 2000);\n'
        f'        {class_name} b2 = new {class_name}("beta", "tls-beta", 2000);\n'
        "        assertTrue(a2.equals(b2));\n"
        "\n"
        f'        {class_name} a3 = new {class_name}("gamma", 3000);\n'
        f'        {class_name} b3 = new {class_name}("delta", 3000);\n'
        "        assertFalse(a3.equals(b3));\n"
        "    }\n"
        "}\n"
    )


# ---------------------------------------------------------------------------
# Demo testgen response  (from testgen.py)
# ---------------------------------------------------------------------------


async def hack_for_demo_java_testgen(data: TestGenSchema) -> TestGenResponseSchema:
    # Extract package and class dynamically from the source code
    source_code = data.source_code_being_tested

    from core.languages.java.testgen import _extract_class_from_source, _extract_package_from_source  # noqa: PLC0415

    package_name = _extract_package_from_source(source_code) or ""
    class_name = _extract_class_from_source(source_code) or data.function_to_optimize.function_name
    test_framework = data.test_framework if data.test_framework in ("junit4", "junit5") else "junit5"

    test_index = data.test_index if data.test_index is not None else 0

    if is_host_equals_demo(source_code):
        if test_index == 0:
            generated_test_source = _build_host_equals_demo_test_source_0(package_name, class_name, test_framework)
        else:
            generated_test_source = _build_host_equals_demo_test_source_1(package_name, class_name, test_framework)
    elif test_index == 0:
        generated_test_source = _build_demo_test_source_0(package_name, class_name, test_framework)
    else:
        generated_test_source = _build_demo_test_source_1(package_name, class_name, test_framework)

    await asyncio.sleep(5)
    # For Java, instrumentation is done client-side
    return TestGenResponseSchema(
        generated_tests=generated_test_source,
        instrumented_behavior_tests=generated_test_source,
        instrumented_perf_tests=generated_test_source,
    )


# ---------------------------------------------------------------------------
# Wrapper functions — single entry points for callers
# ---------------------------------------------------------------------------


async def try_demo_optimize_java(source_code: str) -> OptimizeResponseSchema | None:
    """Check if source_code matches a Java demo function and return pre-canned optimizer results, or None."""
    if not should_hack_for_demo_java(source_code):
        return None
    return await hack_for_demo_java(source_code)


async def try_demo_optimize_java_lp(source_code: str) -> OptimizeResponseSchema | None:
    """Check if source_code matches a Java demo function and return pre-canned LP optimizer results, or None."""
    if not should_hack_for_demo_java(source_code):
        return None
    return await hack_for_demo_java_lp(source_code)


async def try_demo_testgen_java(data: TestGenSchema) -> TestGenResponseSchema | None:
    """Check if data matches a Java demo function and return pre-canned testgen results, or None."""
    if not should_hack_for_demo_java(data.source_code_being_tested):
        return None
    return await hack_for_demo_java_testgen(data)
