from __future__ import annotations

from typing import TYPE_CHECKING, Any

import sentry_sdk
from asgiref.sync import sync_to_async
from django.db import transaction
from ninja import NinjaAPI

from aiservice.common_utils import validate_trace_id
from authapp.auth import AuthenticatedRequest
from log_features.models import LoggingErrorResponseSchema, LoggingSchema, OptimizationFeatures

if TYPE_CHECKING:
    import datetime as dt

features_api = NinjaAPI(urls_namespace="log_features")


@sync_to_async
@transaction.atomic
def log_features(
    trace_id: str,
    user_id: str,
    original_code: str | None = None,
    dependency_code: str | None = None,
    line_profiler_results: str | None = None,
    optimizations_raw: dict[str, str] | None = None,
    optimizations_post: dict[str, str] | None = None,
    explanations_raw: dict[str, str] | None = None,
    explanations_post: dict[str, str] | None = None,
    speedup_ratio: dict[str, float | None] | None = None,
    original_runtime: float | None = None,
    optimized_runtime: dict[str, float] | None = None,
    optimized_line_profiler_results: dict[str, str] | None = None,
    is_correct: dict[str, bool | None] | None = None,
    generated_tests: list[str] | None = None,
    instrumented_generated_tests: list[str] | None = None,
    instrumented_perf_tests: list[str] | None = None,
    test_framework: str | None = None,
    datetime: dt.datetime | None = None,
    aiservice_commit: str | None = None,
    metadata: dict[str, Any] | None = None,
    experiment_metadata: dict[str, str] | None = None,
    final_explanation: str | None = None,
    ranking: dict[str, Any] | None = None,
    optimizations_origin: dict[str, dict[str, str]] | None = None,
) -> None:
    """Log features of a code optimization run to the database.

    :rtype: None
    :param optimized_line_profiler_results: mapping of optimization candidate trace ids to line profiler results
    :param trace_id: The client generated UUID of the optimization run. This is used to link the features together.
    :param user_id: The user ID of the user who ran the optimization.
    :param original_code: The original code that the LLM is allowed to modify.
    :param dependency_code: The dependency code that the LLM is not allowed to modify.
    :param original_code: The line profiling results for the original code.
    :param optimizations_raw: The raw optimizations that were generated by the language model.
    :param optimizations_post: The final optimizations that were generated by the optimization endpoint.
    :param explanations_raw: Raw Explanations generated by the language model.
    :param explanations_post: Final Explanations for the optimizations.
    :param speedup_ratio: Speedups in fractions achieved by the optimizations.
    :param original_runtime: The time taken in ns to run the original code.
    :param optimized_runtime: The time taken in ns to run the optimized code.
    :param is_correct: Behavioural correctness of the optimized code.
    :param generated_tests: Generated tests for the optimized code, output of the aiservice test generation endpoint.
    :param instrumented_generated_tests: Behavior instrumented tests for the optimized code.
    :param instrumented_perf_tests: Performance instrumented tests for the optimized code.
    :param test_framework: The test framework used to generate the tests.
    :param datetime: The datetime of the optimization run. Should be calculated by the aiservice.
    :param aiservice_commit: The commit hash of the AIService code used for the feature logging. hopefully should be the same for the entire run.
    :param metadata: Additional metadata to log.
    :param final_explanation: the final explanation in the PR
    """
    f, created = OptimizationFeatures.objects.select_for_update().get_or_create(
        trace_id=trace_id,
        defaults={
            "user_id": user_id,
            "original_code": original_code,
            "dependency_code": dependency_code,
            "line_profiler_results": line_profiler_results,
            "optimizations_raw": optimizations_raw,
            "optimizations_post": optimizations_post,
            "explanations_raw": explanations_raw,
            "explanations_post": explanations_post,
            "speedup_ratio": speedup_ratio,
            "original_runtime": original_runtime,
            "optimized_runtime": optimized_runtime,
            "optimized_line_profiler_results": optimized_line_profiler_results,
            "is_correct": is_correct,
            "generated_test": generated_tests,
            "instrumented_generated_test": instrumented_generated_tests,
            "instrumented_perf_test": instrumented_perf_tests,
            "test_framework": test_framework,
            "created_at": datetime,
            "aiservice_commit_id": aiservice_commit,
            "metadata": metadata,
            "experiment_metadata": experiment_metadata,
            "final_explanation": final_explanation,
            "ranking": ranking,
            "optimizations_origin": optimizations_origin,
        },
    )

    if created:
        return

    update_fields: list[str] = []

    # Simple override fields - set if new value is provided
    simple_updates = [
        ("user_id", user_id),
        ("original_code", original_code),
        ("dependency_code", dependency_code),
        ("line_profiler_results", line_profiler_results),
        ("final_explanation", final_explanation),
        ("speedup_ratio", speedup_ratio),
        ("original_runtime", original_runtime),
        ("optimized_runtime", optimized_runtime),
        ("optimized_line_profiler_results", optimized_line_profiler_results),
        ("is_correct", is_correct),
        ("test_framework", test_framework),
        ("created_at", datetime),
        ("aiservice_commit_id", aiservice_commit),
    ]
    for field, value in simple_updates:
        if value is not None:
            setattr(f, field, value)
            update_fields.append(field)

    # Dict merge fields - merge new values into existing (only if new value provided)
    dict_merge_updates = [
        ("optimizations_raw", f.optimizations_raw, optimizations_raw),
        ("optimizations_post", f.optimizations_post, optimizations_post),
        ("explanations_raw", f.explanations_raw, explanations_raw),
        ("explanations_post", f.explanations_post, explanations_post),
        ("ranking", f.ranking, ranking),
        ("metadata", f.metadata, metadata),
        ("experiment_metadata", f.experiment_metadata, experiment_metadata),
    ]
    for field, existing, new in dict_merge_updates:
        if new is not None:
            setattr(f, field, merge_dict(existing, new))
            update_fields.append(field)

    # Nested dict merge for optimizations_origin
    if optimizations_origin is not None:
        f.optimizations_origin = merge_dicts(f.optimizations_origin or {}, optimizations_origin)
        update_fields.append("optimizations_origin")

    # List append fields
    if generated_tests:
        f.generated_test = (f.generated_test or []) + generated_tests
        update_fields.append("generated_test")
    if instrumented_generated_tests:
        f.instrumented_generated_test = (f.instrumented_generated_test or []) + instrumented_generated_tests
        update_fields.append("instrumented_generated_test")
    if instrumented_perf_tests:
        f.instrumented_perf_test = (f.instrumented_perf_test or []) + instrumented_perf_tests
        update_fields.append("instrumented_perf_test")

    if update_fields:
        f.save(update_fields=update_fields)


def merge_dict(existing: dict | None, new: dict | None) -> dict | None:
    """Merge new dict into existing, returning existing if new is None."""
    if new is None:
        return existing
    return (existing or {}) | new


def merge_dicts(a: dict[str, dict[str, str]], b: dict[str, dict[str, str]]) -> dict[str, dict[str, str]]:
    result: dict[str, dict[str, str]] = {}

    for key, inner in a.items():
        result[key] = inner.copy()

    for key, inner in b.items():
        if key not in result:
            result[key] = inner.copy()
        else:
            # b overrides a
            result[key].update(inner)

    return result


@features_api.post("/", response={200: None, 500: LoggingErrorResponseSchema})
async def log_features_cli(
    request: AuthenticatedRequest, data: LoggingSchema
) -> int | tuple[int, LoggingErrorResponseSchema]:
    try:
        if hasattr(request, "should_log_features") and request.should_log_features:
            if not validate_trace_id(data.trace_id):
                raise ValueError("Invalid UUID")

            await log_features(
                trace_id=data.trace_id,
                user_id=request.user,
                speedup_ratio=data.speedup_ratio,
                original_runtime=data.original_runtime,
                optimized_runtime=data.optimized_runtime,
                optimized_line_profiler_results=data.optimized_line_profiler_results,
                is_correct=data.is_correct,
                metadata=data.metadata,
                optimizations_post=data.optimizations_post,
            )
    except Exception as e:  # noqa: BLE001
        e.add_note("Error logging features")
        sentry_sdk.capture_exception()
        return 500, LoggingErrorResponseSchema(error=str(e))

    return 200
