"""JavaScript/TypeScript line profiler optimizer module.

This module handles line profiler-guided optimization for JavaScript and TypeScript code.
"""

from __future__ import annotations

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

import sentry_sdk
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam

from aiservice.analytics.posthog import ph
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import OPTIMIZE_MODEL, calculate_llm_cost, call_llm
from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
from core.languages.python.optimizer.config import MAX_OPTIMIZER_LP_CALLS, get_model_distribution
from core.languages.python.optimizer.context_utils.context_helpers import (
    group_code,
    is_multi_context_js,
    is_multi_context_ts,
    split_markdown_code,
)
from core.languages.python.optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm import LLM


# Get the prompts directory
current_dir = Path(__file__).parent
JS_PROMPTS_DIR = current_dir / "prompts" / "optimizer"

# Fallback to original location if prompts haven't been moved
if not JS_PROMPTS_DIR.exists():
    JS_PROMPTS_DIR = Path(__file__).parent.parent.parent / "optimizer" / "prompts" / "javascript"

# Load JavaScript system prompt
JS_SYSTEM_PROMPT = JS_PROMPTS_DIR / "system_prompt.md"
if JS_SYSTEM_PROMPT.exists():
    JS_SYSTEM_PROMPT_TEXT = JS_SYSTEM_PROMPT.read_text()
else:
    # Fallback for backwards compatibility
    JS_SYSTEM_PROMPT_TEXT = ""

# Pattern to extract code blocks from JavaScript LLM response (single file, no file path)
JS_CODE_PATTERN = re.compile(r"```(?:javascript|js|typescript|ts)\s*\n(.*?)```", re.MULTILINE | re.DOTALL)

# Pattern to extract code blocks with file paths (multi-file context)
JS_CODE_WITH_PATH_PATTERN = re.compile(
    r"```(?:javascript|js|typescript|ts):([^\n]+)\n(.*?)```", re.MULTILINE | re.DOTALL
)

# Line profiler context prompt for JavaScript
JS_LINE_PROF_CONTEXT = """
Here are the results of the line profiling of the JavaScript/TypeScript code you will be optimizing.
The profiling data shows:
- Line numbers with execution counts (hits)
- Time spent on each line (in milliseconds)
- Percentage of total time per line

Use this data to identify performance bottlenecks and focus your optimization on the hottest code paths.

{line_profiler_results}
"""


def extract_js_code_and_explanation(content: str, is_multi_file: bool = False) -> tuple[str | dict[str, str], str]:
    """Extract JavaScript code and explanation from LLM response.

    Args:
        content: The raw LLM response content
        is_multi_file: Whether to expect multi-file format

    Returns:
        Tuple of (code, explanation) where code is a string for single file
        or dict[str, str] for multi-file

    """
    if is_multi_file:
        # Extract all code blocks with file paths
        matches = JS_CODE_WITH_PATH_PATTERN.findall(content)
        if matches:
            file_to_code: dict[str, str] = {}
            first_match_pos = content.find("```")
            explanation = content[:first_match_pos].strip() if first_match_pos > 0 else ""

            for file_path, code in matches:
                file_to_code[file_path.strip()] = code.strip()

            return file_to_code, explanation

        # Fall back to single file extraction
        return extract_js_code_and_explanation(content, is_multi_file=False)

    # Single file extraction
    match = JS_CODE_PATTERN.search(content)
    if match:
        code = match.group(1).strip()
        explanation_end = match.start()
        explanation = content[:explanation_end].strip()
        return code, explanation

    return "", content


def normalize_js_code(code: str) -> str:
    """Normalize JavaScript code for comparison."""
    # Remove single-line comments
    code = re.sub(r"//.*$", "", code, flags=re.MULTILINE)
    # Remove multi-line comments
    code = re.sub(r"/\*.*?\*/", "", code, flags=re.DOTALL)
    # Normalize whitespace
    code = " ".join(code.split())
    return code


async def optimize_javascript_code_line_profiler_single(
    user_id: str,
    trace_id: str,
    source_code: str,
    line_profiler_results: str,
    dependency_code: str | None = None,
    optimize_model: LLM = OPTIMIZE_MODEL,
    language_version: str = "ES2022",
    language: str = "javascript",
    call_sequence: int | None = None,
) -> tuple[OptimizeResponseItemSchema | None, float | None, str]:
    """Optimize JavaScript/TypeScript code using LLMs with line profiler guidance."""
    lang_name = "TypeScript" if language == "typescript" else "JavaScript"
    code_block_tag = "typescript" if language == "typescript" else "javascript"
    logging.info(f"/optimize-line-profiler: Optimizing {lang_name} code.")
    debug_log_sensitive_data(f"Optimizing {lang_name} code for user {user_id}:\n{source_code}")

    # Check if source code is multi-file format
    is_multi_file = is_multi_context_ts(source_code) if language == "typescript" else is_multi_context_js(source_code)
    original_file_to_code: dict[str, str] = {}

    if is_multi_file:
        original_file_to_code = split_markdown_code(source_code, language)
        logging.info(
            f"Multi-file context detected with {len(original_file_to_code)} files: {list(original_file_to_code.keys())}"
        )

    # Format system prompt with language version
    system_prompt = JS_SYSTEM_PROMPT_TEXT.format(language_version=language_version)

    # Build user prompt with line profiler results
    if is_multi_file:
        # For multi-file, identify the first file as the target and others as helper context
        file_paths = list(original_file_to_code.keys())
        target_file = file_paths[0] if file_paths else "main file"
        helper_files = file_paths[1:] if len(file_paths) > 1 else []

        # Build multi-file instructions
        helper_notice = ""
        if helper_files:
            helper_list = ", ".join(f"`{f}`" for f in helper_files)
            helper_notice = f"""
HELPER FILES: {helper_list}
These files contain helper functions that the target function uses. You may optimize these as well if needed.
"""

        multi_file_instructions = f"""
The code is provided in a multi-file format. Each file is wrapped in a code block with its path.

TARGET FILE: `{target_file}`
{helper_notice}
Output the optimized code for each file that you modify. Wrap each file's code in:
```{code_block_tag}:<file_path>
<optimized code>
```

You MUST output the target file. You may also output helper files if you optimize them.
"""
        system_prompt = system_prompt + "\n" + multi_file_instructions

        user_prompt = f"""Optimize the following {lang_name} code for better performance.

{JS_LINE_PROF_CONTEXT.format(line_profiler_results=line_profiler_results)}

Here is the code to optimize:
{source_code}
"""
    else:
        user_prompt = f"""Optimize the following {lang_name} code for better performance.

{JS_LINE_PROF_CONTEXT.format(line_profiler_results=line_profiler_results)}

Here is the code to optimize:
```{code_block_tag}
{source_code}
```
"""

    if dependency_code:
        user_prompt = f"Dependencies (read-only):\n```{code_block_tag}\n{dependency_code}\n```\n\n{user_prompt}"

    obs_context: dict = {}
    if call_sequence is not None:
        obs_context["call_sequence"] = call_sequence

    messages: list[ChatCompletionMessageParam] = [
        ChatCompletionSystemMessageParam(role="system", content=system_prompt),
        ChatCompletionUserMessageParam(role="user", content=user_prompt),
    ]

    try:
        output = await call_llm(
            llm=optimize_model,
            messages=messages,
            call_type="line_profiler",
            trace_id=trace_id,
            user_id=user_id,
            python_version=language_version,  # Reusing python_version field for language version
            context=obs_context,
        )
    except Exception as e:
        logging.exception(f"LLM Code Generation error in {lang_name} line profiler optimizer")
        sentry_sdk.capture_exception(e)
        debug_log_sensitive_data(f"Failed to generate code for source:\n{source_code}")
        return None, None, optimize_model.name

    llm_cost = calculate_llm_cost(output.raw_response, optimize_model)

    debug_log_sensitive_data(f"LLM optimization response:\n{output.raw_response.model_dump_json(indent=2)}")

    if output.raw_response.usage is not None:
        ph(
            user_id,
            "aiservice-optimize-line-profiler-openai-usage",
            properties={"model": optimize_model.name, "usage": output.raw_response.usage.json(), "language": language},
        )

    # Extract code and explanation from response
    extracted_code, explanation = extract_js_code_and_explanation(output.content, is_multi_file=is_multi_file)

    if not extracted_code:
        sentry_sdk.capture_message(f"No code block found in {lang_name} line profiler optimization response")
        debug_log_sensitive_data(f"No code found in response for source:\n{source_code}")
        return None, llm_cost, optimize_model.name

    optimization_id = str(uuid.uuid4())

    if is_multi_file and isinstance(extracted_code, dict):
        # Handle multi-file response
        # LLM can optimize both target and helper files
        merged_file_to_code: dict[str, str] = {}
        has_changes = False

        for file_path, original_code in original_file_to_code.items():
            if file_path in extracted_code:
                new_code = extracted_code[file_path]

                # Validate the new code
                if language == "typescript":
                    is_valid, error = validate_typescript_syntax(new_code)
                else:
                    is_valid, error = validate_javascript_syntax(new_code)

                if not is_valid:
                    sentry_sdk.capture_message(f"Invalid {lang_name} generated for {file_path}: {error}")
                    debug_log_sensitive_data(f"Invalid code generated for {file_path}:\n{new_code}\nError: {error}")
                    # Keep original code for this file
                    merged_file_to_code[file_path] = original_code
                else:
                    merged_file_to_code[file_path] = new_code
                    if normalize_js_code(new_code) != normalize_js_code(original_code):
                        has_changes = True
            else:
                # File not in response, keep original
                merged_file_to_code[file_path] = original_code

        if not has_changes:
            debug_log_sensitive_data("Generated code identical to original (multi-file)")
            return None, llm_cost, optimize_model.name

        # Format as multi-file markdown
        wrapped_code = group_code(merged_file_to_code, language=code_block_tag)

        result = OptimizeResponseItemSchema(
            source_code=wrapped_code, explanation=explanation, optimization_id=optimization_id
        )
        return result, llm_cost, optimize_model.name

    # Single file handling
    optimized_code = extracted_code if isinstance(extracted_code, str) else ""

    if not optimized_code:
        return None, llm_cost, optimize_model.name

    # Validate the generated code
    if language == "typescript":
        is_valid, error = validate_typescript_syntax(optimized_code)
    else:
        is_valid, error = validate_javascript_syntax(optimized_code)

    if not is_valid:
        sentry_sdk.capture_message(f"Invalid {lang_name} generated: {error}")
        debug_log_sensitive_data(f"Invalid code generated:\n{optimized_code}\nError: {error}")
        return None, llm_cost, optimize_model.name

    # Check that the code is actually different from the original
    if normalize_js_code(optimized_code) == normalize_js_code(source_code):
        debug_log_sensitive_data("Generated code identical to original")
        return None, llm_cost, optimize_model.name

    # Wrap code in markdown format for CLI parsing
    wrapped_code = (
        f"```{code_block_tag}\n{optimized_code}\n```"
        if not optimized_code.endswith("\n")
        else f"```{code_block_tag}\n{optimized_code}```"
    )
    result = OptimizeResponseItemSchema(
        source_code=wrapped_code, explanation=explanation, optimization_id=optimization_id
    )

    return result, llm_cost, optimize_model.name


async def optimize_javascript_code_line_profiler(
    user_id: str,
    trace_id: str,
    source_code: str,
    line_profiler_results: str,
    dependency_code: str | None = None,
    language_version: str = "ES2022",
    language: str = "javascript",
    n_candidates: int = 0,
) -> tuple[list[OptimizeResponseItemSchema], float, dict[str, str]]:
    """Run parallel JavaScript line profiler optimizations with multiple models."""
    if n_candidates == 0:
        return [], 0.0, {}

    model_distribution = get_model_distribution(n_candidates, MAX_OPTIMIZER_LP_CALLS)
    tasks: list[asyncio.Task[tuple[OptimizeResponseItemSchema | None, float | None, str]]] = []
    call_sequence = 1

    async with asyncio.TaskGroup() as tg:
        for model, num_calls in model_distribution:
            for _ in range(num_calls):
                task = tg.create_task(
                    optimize_javascript_code_line_profiler_single(
                        user_id=user_id,
                        trace_id=trace_id,
                        source_code=source_code,
                        line_profiler_results=line_profiler_results,
                        dependency_code=dependency_code,
                        optimize_model=model,
                        language_version=language_version,
                        language=language,
                        call_sequence=call_sequence,
                    )
                )
                tasks.append(task)
                call_sequence += 1

    # Collect results
    optimization_results: list[OptimizeResponseItemSchema] = []
    total_cost = 0.0
    optimization_models: dict[str, str] = {}

    for task in tasks:
        result, cost, model_name = task.result()
        if cost:
            total_cost += cost
        if result is not None:
            optimization_results.append(result)
            optimization_models[result.optimization_id] = model_name

    return optimization_results, total_cost, optimization_models
