"""Java test generation module.

This module generates JUnit 5 tests for Java functions.
Instrumentation is handled by the codeflash CLI client, not here.
"""

from __future__ import annotations

import logging
import os
import re
from pathlib import Path
from typing import TYPE_CHECKING

import sentry_sdk
from openai.types.chat import ChatCompletionMessageParam

from aiservice.analytics.posthog import ph
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import llm_client
from aiservice.llm_models import EXECUTE_MODEL
from authapp.auth import AuthenticatedRequest
from core.log_features.log_event import update_optimization_cost
from core.shared.testgen_models import (
    TestGenerationFailedError,
    TestGenErrorResponseSchema,
    TestGenResponseSchema,
    TestGenSchema,
)

if TYPE_CHECKING:
    from aiservice.llm_models import LLM

from aiservice.validators.java_validator import validate_java_syntax

_TEST_FUNC_RE = re.compile(r"@Test\s*\n\s*(?:public\s+)?void\s+\w+")

# Get the directory of the current file
current_dir = Path(__file__).parent
JAVA_PROMPTS_DIR = current_dir / "prompts" / "testgen"

# Ensure prompts directory exists
JAVA_PROMPTS_DIR.mkdir(parents=True, exist_ok=True)

# Java system prompt for test generation - JUnit 5 (default)
JAVA_SYSTEM_PROMPT_JUNIT5 = """You are an expert Java developer specializing in writing comprehensive JUnit 5 tests.
Your task is to generate high-quality unit tests for the given Java function.

Guidelines:
1. Use JUnit 5 (Jupiter) test annotations (@Test, @BeforeEach, etc.)
2. Include imports for org.junit.jupiter.api.*
3. Use assertEquals, assertTrue, assertFalse, assertThrows as appropriate
4. Test edge cases, boundary conditions, and typical use cases
5. Use descriptive test method names following the pattern: test<Scenario>_<ExpectedBehavior>
6. Include a test class with the pattern: <ClassName>Test
7. Do NOT use mocks unless absolutely necessary
8. Keep tests simple and focused on one assertion per test when possible

CRITICAL - Java Syntax Rules:
- Java does NOT support import aliasing (e.g., "import X as Y" is INVALID)
- Use fully qualified names or regular imports only
- Example: Use "java.util.Base64.getEncoder()" or "import java.util.Base64;" NOT "import java.util.Base64 as Foo;"

CRITICAL - Handling Complex Parameter Types:
- If a function parameter is an abstract class or interface (e.g., Value, Key), use the REAL factory methods or concrete implementations from the library, NOT custom mock classes
- NEVER create a custom class that extends an abstract class unless you implement ALL abstract methods
- Prefer using static factory methods like Value.get("string"), Value.get(123), etc.
- If the real class has a builder or factory pattern, use it
- Check if there are concrete implementations you can instantiate directly

The target Java version is {language_version}.

Function to test: {function_name}
"""

# Java system prompt for JUnit 4
JAVA_SYSTEM_PROMPT_JUNIT4 = """You are an expert Java developer specializing in writing comprehensive JUnit 4 tests.
Your task is to generate high-quality unit tests for the given Java function.

Guidelines:
1. Use JUnit 4 test annotations (@Test, @Before, etc.) - NOT JUnit 5
2. Include imports for org.junit.* (NOT org.junit.jupiter.api.*)
3. Use static imports: import static org.junit.Assert.*
4. Use assertEquals, assertTrue, assertFalse as appropriate. For exceptions, use @Test(expected=ExceptionType.class)
5. Test edge cases, boundary conditions, and typical use cases
6. Use descriptive test method names following the pattern: test<Scenario>_<ExpectedBehavior>
7. Include a test class with the pattern: <ClassName>Test
8. Do NOT use mocks unless absolutely necessary
9. Keep tests simple and focused on one assertion per test when possible

CRITICAL - Java Syntax Rules:
- Java does NOT support import aliasing (e.g., "import X as Y" is INVALID)
- Use fully qualified names or regular imports only
- Example: Use "java.util.Base64.getEncoder()" or "import java.util.Base64;" NOT "import java.util.Base64 as Foo;"

CRITICAL - Handling Complex Parameter Types:
- If a function parameter is an abstract class or interface (e.g., Value, Key), use the REAL factory methods or concrete implementations from the library, NOT custom mock classes
- NEVER create a custom class that extends an abstract class unless you implement ALL abstract methods
- Prefer using static factory methods like Value.get("string"), Value.get(123), etc.
- If the real class has a builder or factory pattern, use it
- Check if there are concrete implementations you can instantiate directly

The target Java version is {language_version}.

Function to test: {function_name}
"""

JAVA_USER_PROMPT_JUNIT5 = """Generate JUnit 5 tests for the following Java function.

Function name: {function_name}
Class name: {class_name}
Full qualified name: {module_path}
Package: {package_name}

Source code:
```java
{function_code}
```

Generate comprehensive unit tests that cover:
1. Basic functionality with typical inputs
2. Edge cases (empty inputs, null values, boundary conditions)
3. Error conditions (if the function can throw exceptions)
4. Large-scale inputs for performance verification

IMPORTANT REQUIREMENTS:
1. The package declaration MUST be: package {package_name};
2. You MUST import the class under test: import {module_path};
3. The test class MUST be named {class_name}Test
4. Create an instance of {class_name} in @BeforeEach or directly in test methods
5. CRITICAL: Do NOT create custom classes that extend abstract classes or interfaces.
   Instead, use factory methods or concrete implementations from the library.
   Example: For com.aerospike.client.Value, use Value.get("string"), Value.get(123), etc.

Wrap your response in a Java code block (```java ... ```).

Example structure:
```java
package {package_name};

import org.junit.jupiter.api.*;
import static org.junit.jupiter.api.Assertions.*;
import {module_path};

public class {class_name}Test {{
    private {class_name} instance;

    @BeforeEach
    void setUp() {{
        instance = new {class_name}();
    }}

    @Test
    void testBasicFunctionality() {{
        // Test code here
    }}
}}
```
"""

JAVA_USER_PROMPT_JUNIT4 = """Generate JUnit 4 tests for the following Java function.

Function name: {function_name}
Class name: {class_name}
Full qualified name: {module_path}
Package: {package_name}

Source code:
```java
{function_code}
```

Generate comprehensive unit tests that cover:
1. Basic functionality with typical inputs
2. Edge cases (empty inputs, null values, boundary conditions)
3. Error conditions (if the function can throw exceptions)
4. Large-scale inputs for performance verification

IMPORTANT REQUIREMENTS:
1. The package declaration MUST be: package {package_name};
2. You MUST import the class under test: import {module_path};
3. The test class MUST be named {class_name}Test
4. Use JUnit 4 annotations and imports (NOT JUnit 5/Jupiter)
5. Create an instance of {class_name} in @Before or directly in test methods
6. CRITICAL: Do NOT create custom classes that extend abstract classes or interfaces.
   Instead, use factory methods or concrete implementations from the library.
   Example: For com.aerospike.client.Value, use Value.get("string"), Value.get(123), etc.

Wrap your response in a Java code block (```java ... ```).

Example structure:
```java
package {package_name};

import org.junit.Test;
import org.junit.Before;
import static org.junit.Assert.*;
import {module_path};

public class {class_name}Test {{
    private {class_name} instance;

    @Before
    public void setUp() {{
        instance = new {class_name}();
    }}

    @Test
    public void testBasicFunctionality() {{
        // Test code here
    }}
}}
```
"""

# Pattern to extract Java code blocks
JAVA_PATTERN = re.compile(r"^```(?:java)?\s*\n(.*?)\n```", re.MULTILINE | re.DOTALL)


def build_java_prompt(
    function_name: str,
    function_code: str,
    module_path: str,
    class_name: str,
    package_name: str,
    test_framework: str = "junit5",
    language_version: str = "17",
) -> tuple[list[ChatCompletionMessageParam], str]:
    """Build the prompt messages for Java test generation.

    Args:
        function_name: Name of the function to test
        function_code: Source code of the function
        module_path: Import path for the module
        class_name: Name of the class containing the function
        package_name: Package name for the test class
        test_framework: Test framework to use ("junit5" or "junit4")
        language_version: Target Java version (e.g., "11", "17", "21")

    Returns:
        Tuple of (messages, posthog_event_suffix)

    """
    # Select prompts based on test framework
    if test_framework == "junit4":
        system_prompt = JAVA_SYSTEM_PROMPT_JUNIT4
        user_prompt = JAVA_USER_PROMPT_JUNIT4
    else:
        system_prompt = JAVA_SYSTEM_PROMPT_JUNIT5
        user_prompt = JAVA_USER_PROMPT_JUNIT5

    system_message: ChatCompletionMessageParam = {
        "role": "system",
        "content": system_prompt.format(function_name=function_name, language_version=f"Java {language_version}"),
    }

    user_message: ChatCompletionMessageParam = {
        "role": "user",
        "content": user_prompt.format(
            function_name=function_name,
            function_code=function_code,
            module_path=module_path,
            class_name=class_name,
            package_name=package_name,
        ),
    }

    messages: list[ChatCompletionMessageParam] = [system_message, user_message]
    return messages, f"java-{test_framework}-"


def parse_and_validate_java_output(response_content: str) -> str:
    """Parse and validate the LLM response for Java code.

    Args:
        response_content: Raw LLM response

    Returns:
        Validated Java code

    Raises:
        ValueError: If no valid code block found
        SyntaxError: If code has syntax errors

    """
    # Check for code block
    if "```" not in response_content:
        sentry_sdk.capture_message("LLM response did not contain a code block:\n" + response_content[:500])
        raise ValueError("LLM response did not contain a code block.")

    pattern_res = JAVA_PATTERN.search(response_content)
    if not pattern_res:
        raise ValueError("No Java code block found in the LLM response.")

    code = pattern_res.group(1).strip()

    # Syntax validation using tree-sitter
    is_valid, error = validate_java_syntax(code)
    if not is_valid:
        raise SyntaxError(f"Invalid Java code: {error}")

    # Check for test functions
    if not _has_test_functions(code):
        raise ValueError("Generated code does not contain any @Test annotated methods.")

    return code


def _has_test_functions(code: str) -> bool:
    """Check if the code contains JUnit test functions."""
    return _TEST_FUNC_RE.search(code) is not None


async def generate_and_validate_java_test_code(
    messages: list[ChatCompletionMessageParam],
    model: LLM,
    cost_tracker: list[float],
    user_id: str,
    posthog_event_suffix: str,
    trace_id: str = "",
    call_sequence: int | None = None,
    test_index: int | None = None,
) -> str:
    """Generate and validate Java test code using an LLM.

    Args:
        messages: Prompt messages for the LLM
        model: LLM model to use
        cost_tracker: List to track costs
        user_id: User ID for analytics
        posthog_event_suffix: Suffix for analytics events
        trace_id: Trace ID for logging
        call_sequence: Call sequence number
        test_index: Index of the test group being generated

    Returns:
        Validated Java test code

    Raises:
        ValueError: If code generation fails
        SyntaxError: If generated code has syntax errors

    """
    obs_context: dict[str, object] | None = (
        {"call_sequence": call_sequence, "test_index": test_index} if call_sequence is not None else None
    )
    output = await llm_client.call(
        llm=model,
        messages=messages,
        call_type="testgen",
        trace_id=trace_id,
        user_id=user_id,
        python_version="N/A",  # Not applicable for Java
        context=obs_context,
    )

    llm_cost = output.cost
    cost_tracker.append(llm_cost)

    debug_log_sensitive_data(f"LLM testgen response:\n{output.content}")

    return parse_and_validate_java_output(output.content)


def _extract_class_and_package(module_path: str) -> tuple[str, str]:
    """Extract class name and package from module path.

    Args:
        module_path: e.g., "com.example.Algorithms"

    Returns:
        Tuple of (class_name, package_name)

    """
    parts = module_path.rsplit(".", 1)
    if len(parts) == 2:
        return parts[1], parts[0]  # class_name, package_name
    return parts[0], ""  # class_name only, no package


def _extract_package_from_source(source_code: str) -> str | None:
    """Extract package name from Java source code.

    Args:
        source_code: Java source code

    Returns:
        Package name (e.g., "com.example"), or None if not found

    """
    # First try: package declaration in source (most reliable)
    package_pattern = re.compile(r"^\s*package\s+([\w.]+)\s*;", re.MULTILINE)
    match = package_pattern.search(source_code)
    if match:
        logging.debug("Extracted package from declaration: %s", match.group(1))
        return match.group(1)

    # Second try: extract from markdown code block header (e.g., "```java:src/main/java/com/example/Algorithms.java")
    markdown_pattern = re.compile(r"```java:([^\n`]+\.java)", re.IGNORECASE)
    markdown_match = markdown_pattern.search(source_code)
    if markdown_match:
        file_path = markdown_match.group(1).strip()
        package = _extract_package_from_path(file_path)
        if package:
            logging.debug("Extracted package from markdown header: %s", package)
            return package

    # Third try: extract from file path comment (e.g., "// file: src/main/java/com/example/Algorithms.java")
    # Also handle "// file: src/com/example/Algorithms.java" (non-standard Maven)
    file_comment_pattern = re.compile(r"//\s*file:\s*([^\n]+\.java)", re.IGNORECASE)
    file_match = file_comment_pattern.search(source_code)
    if file_match:
        file_path = file_match.group(1).strip()
        logging.debug("Found file comment: %s", file_path)
        package = _extract_package_from_path(file_path)
        if package:
            logging.debug("Extracted package from file comment: %s", package)
            return package

    # Fourth try: infer package from import statements (last resort)
    # Look for imports that might indicate the package structure
    import_pattern = re.compile(r"^\s*import\s+([\w.]+)\.[\w*]+\s*;", re.MULTILINE)
    imports = import_pattern.findall(source_code)
    if imports:
        # Find common prefix among imports that look like internal packages
        # Exclude common library packages
        internal_imports: list[str] = [
            imp for imp in imports if not imp.startswith(("java.", "javax.", "org.junit", "org.apache", "com.google"))
        ]
        if internal_imports:
            # Use the shortest import path as a hint
            internal_imports.sort(key=len)
            logging.debug("Inferred package from imports: %s", internal_imports[0])
            return internal_imports[0]

    logging.warning("Could not extract package name from source code")
    return None


def _extract_package_from_path(file_path: str) -> str | None:
    """Extract Java package from a file path.

    Args:
        file_path: e.g., "src/main/java/com/example/Algorithms.java" or "src/com/example/Algorithms.java"

    Returns:
        Package name (e.g., "com.example"), or None if not found

    """
    # Normalize slashes
    file_path = file_path.replace("\\", "/")

    # Standard Maven patterns (highest priority)
    java_src_patterns = ["/src/main/java/", "/src/test/java/", "src/main/java/", "src/test/java/"]
    for pattern in java_src_patterns:
        if pattern in file_path:
            idx = file_path.find(pattern)
            remaining = file_path[idx + len(pattern) :]
            parts = remaining.split("/")
            if len(parts) > 1:
                package_parts = parts[:-1]  # Remove the filename
                return ".".join(package_parts)

    # Non-standard paths: look for "src/" followed by what looks like a package structure
    # e.g., "src/com/aerospike/client/util/Crypto.java" -> "com.aerospike.client.util"
    # or "client/src/com/aerospike/client/util/Crypto.java"
    src_patterns = ["/src/", "src/"]
    for pattern in src_patterns:
        if pattern in file_path:
            idx = file_path.find(pattern)
            remaining = file_path[idx + len(pattern) :]
            parts = remaining.split("/")
            if len(parts) > 1:
                # Check if first part looks like a package (lowercase, not 'main', 'test', 'java', 'resources')
                first_part = parts[0].lower()
                if first_part not in ("main", "test", "java", "resources") and first_part.isalpha():
                    package_parts = parts[:-1]  # Remove the filename
                    return ".".join(package_parts)

    return None


def _extract_class_from_source(source_code: str) -> str | None:
    """Extract class name from Java source code.

    Args:
        source_code: Java source code

    Returns:
        Class name, or None if not found

    """
    # First try: class declaration in source
    class_pattern = re.compile(r"\bclass\s+(\w+)")
    match = class_pattern.search(source_code)
    if match:
        return match.group(1)

    # Second try: extract from markdown code block header (e.g., "```java:src/main/java/com/example/Algorithms.java")
    markdown_pattern = re.compile(r"```java:([^\n`]+\.java)", re.IGNORECASE)
    markdown_match = markdown_pattern.search(source_code)
    if markdown_match:
        file_path = markdown_match.group(1).strip()
        filename = os.path.basename(file_path)
        if filename.endswith(".java"):
            return filename[:-5]  # Remove .java extension

    # Third try: extract from file path comment (e.g., "// file: src/main/java/com/example/Algorithms.java")
    file_comment_pattern = re.compile(r"//\s*file:\s*([^\n]+\.java)", re.IGNORECASE)
    file_match = file_comment_pattern.search(source_code)
    if file_match:
        file_path = file_match.group(1).strip()
        # Extract class name from file path (e.g., "Algorithms.java" -> "Algorithms")
        filename = os.path.basename(file_path)
        if filename.endswith(".java"):
            return filename[:-5]  # Remove .java extension

    return None


async def testgen_java(
    request: AuthenticatedRequest, data: TestGenSchema
) -> tuple[int, TestGenResponseSchema | TestGenErrorResponseSchema]:
    """Generate Java tests using LLMs."""
    ph(request.user, "aiservice-testgen-java-called")

    # Validate request
    if not data.function_to_optimize:
        return 400, TestGenErrorResponseSchema(error="Invalid function to optimize. It is empty.")
    if not validate_trace_id(data.trace_id):
        return 400, TestGenErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")

    logging.info("/testgen: Generating Java tests...")

    from core.languages.java.demo_hacks import try_demo_testgen_java  # noqa: PLC0415

    demo_result = await try_demo_testgen_java(data)
    if demo_result is not None:
        return 200, demo_result

    try:
        debug_log_sensitive_data(f"Generating Java tests for function {data.function_to_optimize.function_name}")
        logging.info("Generating Java tests for function %s", data.function_to_optimize.function_name)

        # Extract class and package info from source code (more reliable than qualified_name)
        source_code = data.source_code_being_tested
        debug_log_sensitive_data(f"Source code first 200 chars: {source_code[:200]}")

        package_name = _extract_package_from_source(source_code) or ""
        class_name = _extract_class_from_source(source_code) or data.function_to_optimize.function_name

        # Build the full module path (package.ClassName)
        if package_name:
            module_path = f"{package_name}.{class_name}"
        else:
            module_path = class_name

        # Determine test framework (default to junit5 if not specified)
        test_framework = data.test_framework if data.test_framework in ("junit4", "junit5") else "junit5"

        logging.info(
            "Java testgen: package=%s, class=%s, module_path=%s, framework=%s",
            package_name,
            class_name,
            module_path,
            test_framework,
        )
        debug_log_sensitive_data(
            f"Extracted: package={package_name}, class={class_name}, module_path={module_path}, framework={test_framework}"
        )

        # Build prompt
        messages, posthog_event_suffix = build_java_prompt(
            function_name=data.function_to_optimize.function_name,
            function_code=data.source_code_being_tested,
            module_path=module_path,
            class_name=class_name,
            package_name=package_name,
            test_framework=test_framework,
            language_version=data.language_version or "17",
        )

        # Track costs
        cost_tracker: list[float] = []

        # Generate tests
        test_index = data.test_index if data.test_index is not None else 0
        generated_test_code = await generate_and_validate_java_test_code(
            messages=messages,
            model=EXECUTE_MODEL,
            cost_tracker=cost_tracker,
            user_id=request.user,
            posthog_event_suffix=posthog_event_suffix,
            trace_id=data.trace_id,
            call_sequence=data.call_sequence,
            test_index=test_index,
        )

        # Track analytics
        total_cost = sum(cost_tracker)
        ph(
            request.user,
            f"aiservice-testgen-{posthog_event_suffix}success",
            properties={
                "trace_id": data.trace_id,
                "total_cost": total_cost,
                "test_count": len(_TEST_FUNC_RE.findall(generated_test_code)),
            },
        )

        # Update cost tracking
        await update_optimization_cost(data.trace_id, total_cost, request.user)

        # For Java, instrumentation is done client-side
        # Return the generated tests without server-side instrumentation
        return 200, TestGenResponseSchema(
            generated_tests=generated_test_code,
            instrumented_behavior_tests=generated_test_code,  # Client will instrument
            instrumented_perf_tests=generated_test_code,  # Client will instrument
        )

    except TestGenerationFailedError as e:
        logging.warning("Java test generation failed: %s", e)
        sentry_sdk.capture_exception(e)
        return 422, TestGenErrorResponseSchema(error=str(e))
    except (ValueError, SyntaxError) as e:
        logging.warning("Java test generation error: %s", e)
        sentry_sdk.capture_exception(e)
        return 422, TestGenErrorResponseSchema(error=f"Failed to generate valid tests: {e}")
    except Exception as e:
        logging.exception("Unexpected error in Java test generation")
        sentry_sdk.capture_exception(e)
        return 500, TestGenErrorResponseSchema(error=f"Internal error: {e}")
