"""JavaScript/TypeScript test generation module.

This module generates Jest tests for JavaScript/TypeScript functions.
Instrumentation is handled by the codeflash CLI client, not here.
"""

from __future__ import annotations

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

import sentry_sdk
from ninja.errors import HttpError
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, HAIKU_MODEL, OPENAI_MODEL
from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
from authapp.auth import AuthenticatedRequest
from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import safe_log_features
from core.shared.jinja_utils import create_prompt_env
from core.shared.testgen_models import (
    TestGenerationFailedError,
    TestGenErrorResponseSchema,
    TestGenResponseSchema,
    TestGenSchema,
)

if TYPE_CHECKING:
    from aiservice.llm_models import LLM

_JS_RESERVED_WORDS = frozenset(
    {
        "module",
        "exports",
        "require",
        "global",
        "process",  # Node.js globals
        "prototype",
        "constructor",
        "__proto__",  # Object properties
        "break",
        "case",
        "catch",
        "continue",
        "debugger",
        "default",
        "delete",
        "do",
        "else",
        "finally",
        "for",
        "function",
        "if",
        "in",
        "instanceof",
        "new",
        "return",
        "switch",
        "this",
        "throw",
        "try",
        "typeof",
        "var",
        "void",
        "while",
        "with",
        "class",
        "const",
        "enum",
        "export",
        "extends",
        "import",
        "super",
        "implements",
        "interface",
        "let",
        "package",
        "private",
        "protected",
        "public",
        "static",
        "yield",
        "null",
        "true",
        "false",
    }
)

_DEFAULT_CLASS_PATTERN = re.compile(r"\bexport\s+default\s+class\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\b")

_DEFAULT_FUNC_PATTERN = re.compile(r"\bexport\s+default\s+function\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\b")

_NAMED_CLASS_PATTERN = re.compile(r"\bexport\s+class\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\b")

_NAMED_FUNC_PATTERN = re.compile(r"\bexport\s+function\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\b")

_NAMED_CONST_PATTERN = re.compile(r"\bexport\s+const\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\b")

_NAMED_BRACES_PATTERN = re.compile(r"\bexport\s*\{[^}]*\b([a-zA-Z_$][a-zA-Z0-9_$]*)\b[^{]*\}")

_TEST_FUNC_RE = re.compile(r"(?:test|it)\s*\(\s*['\"]")

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

_jinja_env = create_prompt_env(JS_PROMPTS_DIR)

# Pattern to extract JavaScript code blocks
JS_PATTERN = re.compile(r"^```(?:javascript|js|typescript|ts)?\s*\n(.*?)\n```", re.MULTILINE | re.DOTALL)


# JavaScript identifier pattern: starts with letter, underscore, or $,
# followed by letters, digits, underscores, or $
JS_IDENTIFIER_PATTERN = re.compile(r"^[a-zA-Z_$][a-zA-Z0-9_$]*$")


def _is_valid_js_identifier(name: str) -> bool:
    """Check if a name is a valid JavaScript identifier (not a reserved word).

    JavaScript identifiers can start with a letter, underscore, or $,
    followed by letters, digits, underscores, or $. This differs from
    Python identifiers (e.g., '$handler' is valid in JS but not Python).
    """
    return bool(JS_IDENTIFIER_PATTERN.match(name)) and name not in _JS_RESERVED_WORDS


# Patterns to strip file extensions from import paths
# LLMs sometimes add .js extensions to TypeScript imports, which breaks module resolution
_JS_EXTENSION_PATTERN = re.compile(r"""(from\s+['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"])""")
_REQUIRE_EXTENSION_PATTERN = re.compile(
    r"""(require\s*\(\s*['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"]\s*\))"""
)
_JEST_MOCK_EXTENSION_PATTERN = re.compile(
    r"""(jest\.(?:mock|doMock|unmock|requireActual|requireMock)\s*\(\s*['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"])"""
)
_VITEST_MOCK_EXTENSION_PATTERN = re.compile(
    r"""(vi\.(?:mock|doMock|unmock|requireActual|requireMock|importActual|importMock)\s*\(\s*['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"])"""
)


def strip_js_extensions(source: str) -> str:
    """Strip .js/.ts/.tsx/.jsx extensions from relative import paths.

    TypeScript and Jest/Vitest module resolution automatically resolve file extensions,
    so adding them explicitly can cause "Cannot find module" errors when the LLM
    adds incorrect extensions (e.g., .js to a .ts file).

    This function removes extensions from:
    - ES module imports: import { x } from '../path/file.js'
    - CommonJS requires: require('../path/file.js')
    - Jest mocks: jest.mock('../path/file.js')
    - Vitest mocks: vi.mock('../path/file.js')

    Args:
        source: The test source code.

    Returns:
        Source code with extensions stripped from relative import paths.

    """
    source = _JS_EXTENSION_PATTERN.sub(r"\1\2\4", source)
    source = _REQUIRE_EXTENSION_PATTERN.sub(r"\1\2\4", source)
    source = _JEST_MOCK_EXTENSION_PATTERN.sub(r"\1\2\4", source)
    return _VITEST_MOCK_EXTENSION_PATTERN.sub(r"\1\2\4", source)


def _detect_export_style(source_code: str, identifier: str) -> str | None:
    """Detect if an identifier is exported and its export style.

    Args:
        source_code: The source code to analyze
        identifier: The class/function name to look for

    Returns:
        "default" if exported as default, "named" if exported as named, None if not exported

    """
    # Normalize whitespace for easier pattern matching
    normalized = re.sub(r"\s+", " ", source_code)

    # Check for default export: export default class X or export default function X
    for match in _DEFAULT_CLASS_PATTERN.finditer(normalized):
        if match.group(1) == identifier:
            return "default"

    for match in _DEFAULT_FUNC_PATTERN.finditer(normalized):
        if match.group(1) == identifier:
            return "default"

    # Check for named export: export class X, export function X, export const X, export { X }
    for match in _NAMED_CLASS_PATTERN.finditer(normalized):
        if match.group(1) == identifier:
            return "named"

    for match in _NAMED_FUNC_PATTERN.finditer(normalized):
        if match.group(1) == identifier:
            return "named"

    for match in _NAMED_CONST_PATTERN.finditer(normalized):
        if match.group(1) == identifier:
            return "named"

    for match in _NAMED_BRACES_PATTERN.finditer(normalized):
        if match.group(1) == identifier:
            return "named"

    return None


def _resolve_import(function_name: str, module_path: str, source_code: str = "") -> tuple[str, str, str]:
    """Determine import style and binding name for a JS/TS function.

    Analyzes the function name and source code to decide how to import it:
    - 'ClassName.method' with 'export class ClassName' → named import
    - 'ClassName.method' with 'export default class ClassName' → default import
    - 'ClassName.method' with no export → named import (will fail, surfacing the issue)
    - 'funcName' → named import
    - Complex patterns → namespace import (fallback)

    The actual ESM vs CJS formatting is handled by the Jinja2 js_import macro.

    Args:
        function_name: Name of the function to test (e.g., "ClassName.method")
        module_path: Import path for the module
        source_code: Source code to analyze for export detection (optional)

    Returns:
        Tuple of (import_style, import_name, function_accessor) where
        import_style is "default", "named", or "namespace".

    """
    parts = function_name.split(".")

    # Handle ClassName.method pattern
    if len(parts) == 2:
        class_name, method_name = parts
        if _is_valid_js_identifier(class_name) and method_name.isidentifier():
            # Detect export style from source code if provided
            if source_code:
                export_style = _detect_export_style(source_code, class_name)
                if export_style == "default":
                    return "default", class_name, function_name
                # For both "named" and None (not exported), use named import
                # If not exported, the test will fail, surfacing the issue that the class needs to be exported
                return "named", class_name, function_name
            # Fallback to named import if no source code provided (safer than assuming default)
            return "named", class_name, function_name

    # Handle standalone function
    if len(parts) == 1 and parts[0].isidentifier():
        return "named", function_name, function_name

    # Fallback to namespace import for complex patterns
    module_name = module_path.rstrip("/").split("/")[-1].replace("-", "_").replace(".", "_")
    if not module_name.isidentifier():
        module_name = "mod"
    return "namespace", module_name, function_name


def build_javascript_prompt(
    function_name: str,
    function_code: str,
    module_path: str,
    test_framework: str,
    is_async: bool,
    module_system: str | None = None,
    model_type: str = "openai",
) -> tuple[list[ChatCompletionMessageParam], str]:
    """Build the prompt messages for JavaScript 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
        test_framework: Testing framework (jest, vitest, mocha)
        is_async: Whether the function is async
        module_system: "esm" or "commonjs" (defaults to CJS when None)
        model_type: "openai" or "anthropic" — controls prompt formatting

    Returns:
        Tuple of (messages, posthog_event_suffix)

    """
    async_prefix = "async_" if is_async else ""
    framework = "mocha_" if test_framework == "mocha" else ""
    system_template = f"{async_prefix}{framework}system.md.j2"
    posthog_event_suffix = "async-" if is_async else ""

    import_style, import_name, function_accessor = _resolve_import(function_name, module_path, function_code)

    system_message: ChatCompletionMessageParam = {
        "role": "system",
        "content": _jinja_env.get_template(system_template).render(
            function_name=function_accessor, model_type=model_type
        ),
    }

    user_message: ChatCompletionMessageParam = {
        "role": "user",
        "content": _jinja_env.get_template("user.md.j2").render(
            is_async=is_async,
            module_system=module_system,
            import_style=import_style,
            import_name=import_name,
            test_framework=test_framework,
            function_name=function_accessor,
            function_code=function_code,
            module_path=module_path,
            package_comment="",
        ),
    }

    messages: list[ChatCompletionMessageParam] = [system_message, user_message]
    return messages, posthog_event_suffix


def parse_and_validate_js_output(response_content: str, language: str = "javascript") -> str:
    """Parse and validate the LLM response for JavaScript/TypeScript code.

    Args:
        response_content: Raw LLM response
        language: Language to validate against ("javascript" or "typescript")

    Returns:
        Validated JavaScript/TypeScript 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 = JS_PATTERN.search(response_content)
    if not pattern_res:
        raise ValueError("No JavaScript code block found in the LLM response.")

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

    # Validate syntax using the appropriate validator based on language
    if language == "typescript":
        is_valid, error = validate_typescript_syntax(code)
        lang_name = "TypeScript"
    else:
        is_valid, error = validate_javascript_syntax(code)
        lang_name = "JavaScript"

    if not is_valid:
        raise SyntaxError(f"Invalid {lang_name} syntax: {error}")

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

    return code


def _has_test_functions(code: str) -> bool:
    """Check if the code contains Jest test functions."""
    # Look for test() or it() calls
    return _TEST_FUNC_RE.search(code) is not None


async def generate_and_validate_js_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,
    language: str = "javascript",
    test_index: int | None = None,
) -> str:
    """Generate and validate JavaScript/TypeScript test code using LLM.

    Args:
        messages: Prompt messages
        model: LLM model to use
        cost_tracker: List to track costs
        user_id: User ID
        posthog_event_suffix: Suffix for PostHog events
        trace_id: Trace ID for logging
        call_sequence: Call sequence number
        language: Language to validate against ("javascript" or "typescript")
        test_index: Index of the test group being generated

    Returns:
        Validated JavaScript/TypeScript test code

    Raises:
        SyntaxError: If code is invalid
        ValueError: If no valid code found

    """
    obs_context: dict[str, int | None] | None = (
        {"call_sequence": call_sequence, "test_index": test_index} if call_sequence is not None else None
    )

    response = await llm_client.call(
        llm=model,
        messages=messages,
        call_type="test_generation",
        trace_id=trace_id,
        user_id=user_id,
        python_version="javascript",  # Reusing field for language
        context=obs_context,
    )

    cost = response.cost
    cost_tracker.append(cost)

    debug_log_sensitive_data(
        f"JavaScript {posthog_event_suffix}execute response:\n{response.raw_response.model_dump_json(indent=2)}"
    )

    if response.raw_response.usage:
        ph(
            user_id,
            f"aiservice-testgen-js-{posthog_event_suffix}execute-openai-usage",
            properties={"model": model.name, "usage": response.raw_response.usage.model_dump_json()},
        )

    # Parse and validate using the appropriate language validator
    validated_code = parse_and_validate_js_output(response.content, language=language)
    return validated_code


async def generate_javascript_tests_from_function(
    user_id: str,
    function_name: str,
    function_code: str,
    module_path: str,
    test_framework: str = "jest",
    execute_model: LLM = EXECUTE_MODEL,
    is_async: bool = False,
    trace_id: str = "",
    call_sequence: int | None = None,
    language: str = "javascript",
    test_index: int | None = None,
    module_system: str | None = None,
    model_type: str = "openai",
) -> tuple[str, str, str]:
    """Generate JavaScript tests for a function.

    Args:
        user_id: User ID
        function_name: Name of function to test
        function_code: Source code of function
        module_path: Import path for module
        test_framework: Testing framework (jest, vitest, mocha)
        execute_model: LLM model to use
        is_async: Whether function is async
        trace_id: Trace ID for logging
        call_sequence: Call sequence number
        language: Language of the function to test (javascript, typescript)
        test_index: Index of the test group being generated
        module_system: "esm" or "commonjs" (defaults to CJS when None)
        model_type: "openai" or "anthropic" — controls prompt formatting

    Returns:
        Tuple of (generated_tests, instrumented_behavior_tests, instrumented_perf_tests)

    Raises:
        TestGenerationFailedError: If test generation fails

    """
    messages, posthog_event_suffix = build_javascript_prompt(
        function_name=function_name,
        function_code=function_code,
        module_path=module_path,
        test_framework=test_framework,
        is_async=is_async,
        module_system=module_system,
        model_type=model_type,
    )

    cost_tracker: list[float] = []

    try:
        validated_code = await generate_and_validate_js_test_code(
            messages=messages,
            model=execute_model,
            cost_tracker=cost_tracker,
            user_id=user_id,
            posthog_event_suffix=posthog_event_suffix,
            trace_id=trace_id,
            call_sequence=call_sequence,
            language=language,
            test_index=test_index,
        )

        total_llm_cost = sum(cost_tracker)
        await update_optimization_cost(trace_id=trace_id, cost=total_llm_cost, user_id=user_id)

        # Return uninstrumented code - instrumentation is handled by the codeflash CLI client
        # This ensures consistent instrumentation using the codeflash npm package
        return validated_code, validated_code, validated_code

    except (SyntaxError, ValueError) as e:
        total_llm_cost = sum(cost_tracker)
        await update_optimization_cost(trace_id=trace_id, cost=total_llm_cost, user_id=user_id)
        msg = f"Failed to generate valid JavaScript test code after {len(cost_tracker)} tries. trace_id={trace_id}"
        logging.exception(msg)
        raise TestGenerationFailedError(msg) from e


def validate_javascript_testgen_request_data(data: TestGenSchema) -> None:
    """Validate JavaScript/TypeScript test generation request data.

    Args:
        data: Request data

    Raises:
        HttpError: If validation fails

    """
    if data.test_framework not in ["jest", "vitest", "mocha"]:
        raise HttpError(400, "Invalid test framework for JavaScript/TypeScript. Supported: jest, vitest, mocha.")
    if not data.function_to_optimize:
        raise HttpError(400, "Invalid function to optimize. It is empty.")
    if not validate_trace_id(data.trace_id):
        raise HttpError(400, "Invalid trace ID. Please provide a valid UUIDv4.")

    # Validate syntax based on language
    if data.language == "typescript":
        is_valid, error = validate_typescript_syntax(data.source_code_being_tested)
        lang_name = "TypeScript"
    else:
        is_valid, error = validate_javascript_syntax(data.source_code_being_tested)
        lang_name = "JavaScript"

    if not is_valid:
        raise HttpError(400, f"Invalid source code. It is not valid {lang_name}: {error}")


async def testgen_javascript(
    request: AuthenticatedRequest, data: TestGenSchema
) -> tuple[int, TestGenResponseSchema | TestGenErrorResponseSchema]:
    """Main endpoint handler for JavaScript test generation.

    Args:
        request: Authenticated request
        data: Test generation request data

    Returns:
        Tuple of (status_code, response)

    """
    language = data.language
    ph(request.user, "aiservice-testgen-called", properties={"language": language})

    try:
        validate_javascript_testgen_request_data(data)
    except HttpError as e:
        e.add_note(f"JavaScript testgen request validation error: {e.status_code} {e.message}")
        sentry_sdk.capture_exception(e)
        return e.status_code, TestGenErrorResponseSchema(error=e.message)

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

    try:
        debug_log_sensitive_data(f"Generating JavaScript tests for function {data.function_to_optimize.function_name}")

        # Using different LLMs for different test_index values
        test_index = data.test_index if data.test_index is not None else 0
        if test_index % 2 == 0:
            execute_model = OPENAI_MODEL
            model_source = "OpenAI"
        else:
            execute_model = HAIKU_MODEL
            model_source = "Anthropic"

        logging.info("Using %s model (%s) for JavaScript test_index %s", model_source, execute_model.name, test_index)

        (
            generated_test_source,
            instrumented_behavior_tests,
            instrumented_perf_tests,
        ) = await generate_javascript_tests_from_function(
            user_id=request.user,
            function_name=data.function_to_optimize.qualified_name,
            function_code=data.source_code_being_tested,
            module_path=data.module_path,
            test_framework=data.test_framework,
            is_async=data.is_async or False,
            trace_id=data.trace_id,
            call_sequence=data.call_sequence,
            execute_model=execute_model,
            language=data.language,
            test_index=test_index,
            module_system=data.module_system,
            model_type=execute_model.model_type,
        )

        # Strip incorrect file extensions from import paths (LLMs sometimes add .js to .ts imports)
        # Must strip from ALL three test outputs since CLI uses instrumented versions
        generated_test_source = strip_js_extensions(generated_test_source)
        instrumented_behavior_tests = strip_js_extensions(instrumented_behavior_tests)
        instrumented_perf_tests = strip_js_extensions(instrumented_perf_tests)

        ph(request.user, "aiservice-testgen-tests-generated", properties={"language": language})

        if hasattr(request, "should_log_features") and request.should_log_features:
            await safe_log_features(
                trace_id=data.trace_id,
                user_id=request.user,
                generated_tests=[generated_test_source],
                instrumented_generated_tests=[instrumented_behavior_tests],
                test_index=test_index,
                test_framework=data.test_framework,
                metadata={
                    "test_timeout": data.test_timeout,
                    "function_to_optimize": data.function_to_optimize.qualified_name,
                    "file_path": data.function_to_optimize.file_path,
                    "language": language,
                },
            )

        return 200, TestGenResponseSchema(
            generated_tests=generated_test_source,
            instrumented_behavior_tests=instrumented_behavior_tests,
            instrumented_perf_tests=instrumented_perf_tests,
        )

    except TestGenerationFailedError as e:
        logging.warning("JavaScript test generation failed: %s trace_id=%s", e, data.trace_id)
        sentry_sdk.capture_exception(e)
        return 422, TestGenErrorResponseSchema(error=str(e))
    except Exception as e:
        logging.exception("JavaScript test generation failed. trace_id=%s", data.trace_id)
        sentry_sdk.capture_exception(e)
        return 500, TestGenErrorResponseSchema(error="Error generating JavaScript tests. Internal server error.")
