from __future__ import annotations

from typing import TYPE_CHECKING, Any, cast

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 core.log_features.models import LoggingErrorResponseSchema, LoggingSchema, OptimizationFeatures

if TYPE_CHECKING:
    import datetime as dt

features_api = NinjaAPI(urls_namespace="log_features")


def _positional_list(items: list[str] | None, index: int | None) -> list[str | None] | None:
    """Place a single item at the correct index, padding with None."""
    if not items or index is None or len(items) != 1:
        return cast("list[str | None] | None", items)
    result: list[str | None] = [None] * (index + 1)
    result[index] = items[0]
    return result


@sync_to_async(thread_sensitive=False)
@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_index: int | 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:
    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": _positional_list(generated_tests, test_index),
            "instrumented_generated_test": _positional_list(instrumented_generated_tests, test_index),
            "instrumented_perf_test": _positional_list(instrumented_perf_tests, test_index),
            "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_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_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)

    if optimizations_origin is not None:
        f.optimizations_origin = merge_dicts(f.optimizations_origin or {}, optimizations_origin)
        update_fields.append("optimizations_origin")

    list_updates = [
        ("generated_test", f.generated_test, generated_tests),
        ("instrumented_generated_test", f.instrumented_generated_test, instrumented_generated_tests),
        ("instrumented_perf_test", f.instrumented_perf_test, instrumented_perf_tests),
    ]
    for field, existing, new_items in list_updates:
        if not new_items:
            continue
        current = existing or []
        if test_index is not None and len(new_items) == 1:
            while len(current) <= test_index:
                current.append(None)
            current[test_index] = new_items[0]
        else:
            current = current + new_items
        setattr(f, field, current)
        update_fields.append(field)

    if update_fields:
        f.save(update_fields=update_fields)


async def safe_log_features(**kwargs: Any) -> None:  # noqa: ANN401
    try:
        await log_features(**kwargs)
    except Exception:  # noqa: BLE001
        sentry_sdk.capture_exception()


def merge_dict(existing: dict | None, new: dict | None) -> dict | 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:
            result[key].update(inner)

    return result


@features_api.post("/", response={200: None, 400: LoggingErrorResponseSchema, 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):
                return 400, LoggingErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")

            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
