from __future__ import annotations

import asyncio
import json
import logging
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Any

import sentry_sdk
import stamina
from ninja import NinjaAPI, Schema
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
from packaging import version

from aiservice.analytics.posthog import ph
from aiservice.common.markdown_utils import extract_code_block_with_context, wrap_code_in_markdown
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import LLMOutputUnparseable, llm_client
from aiservice.llm_models import OPTIMIZATION_REVIEW_MODEL
from authapp.auth import AuthenticatedRequest
from core.log_features.log_event import update_optimization_cost, update_optimization_features_review

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm_models import LLM

optimization_review_api = NinjaAPI(urls_namespace="optimization_review")

current_dir = Path(__file__).parent
SYSTEM_PROMPT = (current_dir / "prompts" / "optimization_review_system_prompt.md").read_text()
USER_PROMPT_TEMPLATE = (current_dir / "prompts" / "user_prompt_template.md").read_text()


class OptimizationReviewErrorSchema(Schema):
    error: str


class ReviewLevel(str, Enum):
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class OptimizationReviewResponseSchema(Schema):
    review: ReviewLevel
    review_explanation: str
    performance_score: int | None = None  # 1-10
    code_quality_score: int | None = None  # 1-10
    hot_path_score: int | None = None  # 1-10


class OptimizationReviewSchema(Schema):
    trace_id: str
    code_diff: str | None = None  # LEGACY: unified diff format (for CLI backward compatibility)
    original_code: str | None = None
    optimized_code: str | None = None
    original_runtime: str
    optimized_runtime: str
    speedup: str
    existing_tests: str | None = None
    generated_tests: str | None = None
    replay_tests: str | None = None
    benchmark_details: str | None = None
    coverage_message: str
    loop_count: int
    explanation: str
    calling_fn_details: str
    language: str = "python"  # Programming language: python, javascript, typescript
    python_version: str | None = None  # For Python only
    codeflash_version: str = "0.18.2"
    call_sequence: int | None = None
    acceptance_reason: str | None = None  # "runtime", "throughput", "concurrency"
    original_concurrency_ratio: float | None = None
    optimized_concurrency_ratio: float | None = None
    original_throughput: int | None = None
    optimized_throughput: int | None = None


def _build_optimization_review_messages(data: OptimizationReviewSchema) -> list[ChatCompletionMessageParam]:
    if version.parse(data.codeflash_version) <= version.parse("0.18.2") and data.generated_tests:
        data.generated_tests = wrap_code_in_markdown(data.generated_tests)

    system_message = ChatCompletionSystemMessageParam(role="system", content=SYSTEM_PROMPT)
    if data.original_code and data.optimized_code:
        # Determine language version string
        if data.language == "python":
            language_version = data.python_version or "Not specified"
        else:
            language_version = "Not specified"

        user_prompt = USER_PROMPT_TEMPLATE.format(
            original_runtime=data.original_runtime,
            optimized_runtime=data.optimized_runtime,
            speedup=data.speedup,
            loop_count=data.loop_count,
            coverage_message=data.coverage_message,
            language=data.language,
            language_version=language_version,
            original_code=data.original_code,
            optimized_code=data.optimized_code,
            explanation=data.explanation,
            generated_tests=data.generated_tests or "Not Available",
            existing_tests=data.existing_tests or "Not Available",
            replay_tests=data.replay_tests or "Not Available",
            benchmark_details=data.benchmark_details or "Not Available",
            calling_fn_details=data.calling_fn_details or "Not Available",
            acceptance_reason=data.acceptance_reason or "runtime",
            original_concurrency_ratio=data.original_concurrency_ratio or "N/A",
            optimized_concurrency_ratio=data.optimized_concurrency_ratio or "N/A",
            original_throughput=data.original_throughput or "N/A",
            optimized_throughput=data.optimized_throughput or "N/A",
        )

    # LEGACY FORMAT: Unified diff (for backward compatibility with old CLI)
    elif data.code_diff:
        # Old CLI sends unified diff - build prompt with diff instead
        language_version = data.python_version or "Not specified" if data.language == "python" else "Not specified"
        user_prompt = f"""
<optimization_review_request>

<performance_metrics>
<original_runtime>{data.original_runtime}</original_runtime>
<optimized_runtime>{data.optimized_runtime}</optimized_runtime>
<speedup_percentage>{data.speedup}</speedup_percentage>
<loop_count>{data.loop_count}</loop_count>
<test_coverage>{data.coverage_message}</test_coverage>
<language>{data.language}</language>
<language_version>{language_version}</language_version>
</performance_metrics>

<code_changes>
{data.code_diff}
</code_changes>

<explanation>
{data.explanation}
</explanation>

<validation_data>
<generated_tests>
{data.generated_tests or "Not Available"}
</generated_tests>

<existing_tests>
{data.existing_tests or "Not Available"}
</existing_tests>

<replay_tests>
{data.replay_tests or "Not Available"}
</replay_tests>

<benchmark_details>
{data.benchmark_details or "Not Available"}
</benchmark_details>
</validation_data>

<codebase_context>
<calling_functions>
{data.calling_fn_details or "Not Available"}
</calling_functions>
</codebase_context>

</optimization_review_request>

Please analyze this optimization following the structured assessment process and provide your rating.
"""
    else:
        raise ValueError("Must provide either (original_code + optimized_code) or code_diff")

    user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)

    return [system_message, user_message]


@stamina.retry(on=LLMOutputUnparseable, attempts=2)
async def get_optimization_review(
    request: AuthenticatedRequest,
    data: OptimizationReviewSchema,
    optimization_review_model: LLM = OPTIMIZATION_REVIEW_MODEL,
) -> tuple[int, OptimizationReviewResponseSchema | OptimizationReviewErrorSchema, float]:
    """Compute optimization review via Claude.

    Returns (status_code, response, llm_cost).
    """
    ph(request.user, "aiservice-optimization-review-called")

    try:
        messages = _build_optimization_review_messages(data)

        debug_log_sensitive_data(f"{messages[0]}{messages[1]}")

        obs_context: dict[str, Any] = {"speedup": data.speedup}
        if data.call_sequence is not None:
            obs_context["call_sequence"] = data.call_sequence

        response = await llm_client.call(
            llm=optimization_review_model,
            messages=messages,
            call_type="optimization_review",
            trace_id=data.trace_id,
            user_id=request.user,
            context=obs_context,
        )

        cost = response.cost

        review_text = response.content.strip()
        if result := extract_code_block_with_context(review_text, language="json"):
            before_text, json_content, after_text = result
            try:
                review_data = json.loads(json_content.lower().strip())
                review_explanation = before_text + " " + after_text

                # Extract scores if present
                performance_score = review_data.get("performance_score")
                code_quality_score = review_data.get("code_quality_score")
                hot_path_score = review_data.get("hot_path_score")

                # Extract reasoning
                performance_reasoning = review_data.get("performance_reasoning", "")
                code_quality_reasoning = review_data.get("code_quality_reasoning", "")
                hot_path_reasoning = review_data.get("hot_path_reasoning", "")

                # Append scores and reasoning to explanation for internal visibility
                if any([performance_score, code_quality_score, hot_path_score]):
                    scores_summary = "\n\n---\n**Dimension Scores:**"
                    if performance_score:
                        scores_summary += f"\n- **Performance: {performance_score}/10**"
                        if performance_reasoning:
                            scores_summary += f"\n  {performance_reasoning}"
                    if code_quality_score:
                        scores_summary += f"\n- **Code Quality: {code_quality_score}/10**"
                        if code_quality_reasoning:
                            scores_summary += f"\n  {code_quality_reasoning}"
                    if hot_path_score:
                        scores_summary += f"\n- **Hot Path: {hot_path_score}/10**"
                        if hot_path_reasoning:
                            scores_summary += f"\n  {hot_path_reasoning}"
                    review_explanation += scores_summary

                review = OptimizationReviewResponseSchema(
                    review=review_data["rating"],
                    review_explanation=review_explanation,
                    performance_score=performance_score,
                    code_quality_score=code_quality_score,
                    hot_path_score=hot_path_score,
                )
            except Exception as e:
                # invalid response
                logging.exception("Invalid optimization review response")
                sentry_sdk.capture_exception(e)
                debug_log_sensitive_data(f"Invalid response : {e}")
                raise LLMOutputUnparseable("Invalid response", cost=cost) from e
            else:
                ph(request.user, "aiservice-optimization-review-successful")
                return 200, review, cost
        else:
            raise LLMOutputUnparseable("Invalid response", cost=cost)
    except LLMOutputUnparseable:
        raise
    except Exception as e:
        logging.exception("Error in optimization_review")
        sentry_sdk.capture_exception(e)
        return 500, OptimizationReviewErrorSchema(error="Internal server error"), 0.0


@optimization_review_api.post(
    "/",
    response={
        200: OptimizationReviewResponseSchema,
        400: OptimizationReviewErrorSchema,
        422: OptimizationReviewErrorSchema,
        500: OptimizationReviewErrorSchema,
    },
)
async def optimization_review(
    request: AuthenticatedRequest, data: OptimizationReviewSchema
) -> tuple[int, OptimizationReviewResponseSchema | OptimizationReviewErrorSchema]:
    if not validate_trace_id(data.trace_id):
        return 400, OptimizationReviewErrorSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    try:
        response_code, output, llm_cost = await get_optimization_review(request, data)
    except LLMOutputUnparseable:
        return 422, OptimizationReviewErrorSchema(error="Invalid response")
    if isinstance(output, OptimizationReviewResponseSchema):
        review_event = output.review.value
        review_explanation = output.review_explanation
    else:
        review_event = output.error
        review_explanation = ""
    results = await asyncio.gather(
        update_optimization_cost(data.trace_id, llm_cost, user_id=request.user),
        update_optimization_features_review(
            trace_id=data.trace_id,
            review_quality=review_event,
            review_explanation=review_explanation,
            calling_fn_details=data.calling_fn_details,
        ),
        return_exceptions=True,
    )
    for result in results:
        if isinstance(result, BaseException):
            sentry_sdk.capture_exception(result)
    return response_code, output
