from __future__ import annotations

import datetime as dt
import logging
import uuid
from asyncio import Semaphore
from typing import Dict, List, Optional

from authapp.auth import AuthBearer
from ninja import NinjaAPI, Schema

from log_features.models import OptimizationFeatures

features_api = NinjaAPI(auth=AuthBearer(), urls_namespace="log_features")


semaphores: dict[str, Semaphore] = {}


def get_semaphore(trace_id: str) -> Semaphore:
    logging.debug("Getting semaphore for trace_id: %s", trace_id)
    if trace_id not in semaphores:
        semaphores[trace_id] = Semaphore(1)
    return semaphores[trace_id]


async def log_features(
    trace_id: str,
    user_id: str,
    original_code: Optional[str] = None,
    optimizations_raw: Optional[Dict[str, str]] = None,
    optimizations_post: Optional[Dict[str, str]] = None,
    explanations_raw: Optional[Dict[str, str]] = None,
    explanations_post: Optional[Dict[str, str]] = None,
    speedup_ratio: Optional[Dict[str, float | None]] = None,
    original_runtime: Optional[float] = None,
    optimized_runtime: Optional[Dict[str, float]] = None,
    is_correct: Optional[Dict[str, bool | None]] = None,
    generated_tests: Optional[List[str]] = None,
    instrumented_generated_tests: Optional[List[str]] = None,
    test_framework: Optional[str] = None,
    datetime: Optional[dt.datetime] = None,
    aiservice_commit: Optional[str] = None,
    metadata: Optional[Dict[str, any]] = None,
    experiment_metadata: Optional[Dict[str, str]] = None,
) -> None:
    """Log features of a code optimization run to the database.

    :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 was optimized. Includes dependencies.
    :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: Instrumented generated tests for the optimized code, output of the aiservice test generation endpoint.
    :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.
    """
    try:
        uuid_obj = uuid.UUID(trace_id, version=4)
        if str(uuid_obj) != trace_id:
            raise ValueError("Invalid UUID")
        try:
            semaphore = get_semaphore(trace_id)
            async with semaphore:
                # look up existing feature if it exists
                f: OptimizationFeatures = await OptimizationFeatures.objects.aget(trace_id=trace_id)
                updated_metadata = f.metadata or {}
                if metadata:
                    updated_metadata.update(metadata)
                updated_experiment_metadata = f.experiment_metadata or {}
                if experiment_metadata:
                    updated_experiment_metadata.update(experiment_metadata)

                if generated_tests:
                    f.generated_test = (f.generated_test or []) + generated_tests
                if instrumented_generated_tests:
                    f.instrumented_generated_test = (f.instrumented_generated_test or []) + instrumented_generated_tests

                # Update fields on the existing instance
                f.user_id = user_id if user_id is not None else f.user_id
                f.original_code = original_code if original_code is not None else f.original_code
                f.optimizations_raw = optimizations_raw if optimizations_raw is not None else f.optimizations_raw
                f.optimizations_post = optimizations_post if optimizations_post is not None else f.optimizations_post
                f.explanations_raw = explanations_raw if explanations_raw is not None else f.explanations_raw
                f.explanations_post = explanations_post if explanations_post is not None else f.explanations_post
                f.speedup_ratio = speedup_ratio if speedup_ratio is not None else f.speedup_ratio
                f.original_runtime = original_runtime if original_runtime is not None else f.original_runtime
                f.optimized_runtime = optimized_runtime if optimized_runtime is not None else f.optimized_runtime
                f.is_correct = is_correct if is_correct is not None else f.is_correct
                f.generated_test = generated_tests if generated_tests is not None else f.generated_test

                f.test_framework = test_framework if test_framework is not None else f.test_framework
                f.created_at = datetime if datetime is not None else f.created_at
                f.aiservice_commit_id = aiservice_commit if aiservice_commit is not None else f.aiservice_commit_id
                f.metadata = updated_metadata
                f.experiment_metadata = updated_experiment_metadata
                await f.asave()

        except OptimizationFeatures.DoesNotExist:
            # no existing row found, create a new one
            await OptimizationFeatures.objects.acreate(
                trace_id=trace_id,
                user_id=user_id,
                original_code=original_code,
                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,
                is_correct=is_correct,
                generated_test=generated_tests,
                instrumented_generated_test=instrumented_generated_tests,
                test_framework=test_framework,
                created_at=datetime,
                aiservice_commit_id=aiservice_commit,
                metadata=metadata,
                experiment_metadata=experiment_metadata,
            )
        finally:
            semaphores.pop(trace_id, None)
    except Exception as e:
        logging.exception(f"Error logging features: {e}")
        raise e


class LoggingSchema(Schema):
    trace_id: str
    speedup_ratio: dict[str, float | int | None] | None
    original_runtime: float | int | None
    optimized_runtime: dict[str, float | int | None] | None
    is_correct: dict[str, bool] | None


class LoggingErrorResponseSchema(Schema):
    error: str


@features_api.post("/", response={200: None, 500: LoggingErrorResponseSchema})
async def log_features_cli(request, data: LoggingSchema):
    try:
        if request.tier is None:
            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,
                is_correct=data.is_correct,
            )
    except Exception as e:
        return 500, LoggingErrorResponseSchema(error=str(e))

    return 200
