from datetime import UTC, datetime
from typing import Any
from uuid import UUID, uuid4

from django.db.models import F
from django.db.models.functions import Coalesce

from aiservice.common_utils import normalize_trace_id
from core.log_features.models import OptimizationEvents, OptimizationFeatures, Repositories


async def _increment_llm_cost(trace_id: str, cost: float) -> None:
    """Atomically increment llm_cost for the given trace_id."""
    await OptimizationEvents.objects.filter(trace_id=trace_id).aupdate(llm_cost=Coalesce(F("llm_cost"), 0.0) + cost)


async def get_repository(owner: str | None, repo_name: str | None) -> Repositories | None:
    if not owner or not repo_name:
        return None
    try:
        full_repo_name = f"{owner}/{repo_name}"
        return await Repositories.objects.aget(full_name=full_repo_name)
    except Repositories.DoesNotExist:
        return None


async def log_optimization_event(
    event_type: str,
    user_id: str | None = None,
    repository_id: str | None = None,
    repo_owner: str | None = None,
    repo_name: str | None = None,
    trace_id: str | None = None,
    api_key_id: int | None = None,
    metadata: dict[str, Any] | None = None,
    current_username: str | None = None,
    llm_cost: float | None = None,
) -> OptimizationEvents:
    if repository_id is None and repo_owner and repo_name:
        repository = await get_repository(repo_owner, repo_name)
        repository_id = str(repository.id) if repository else None

    return await OptimizationEvents.objects.acreate(
        id=str(uuid4()),
        event_type=event_type,
        user_id=user_id,
        repository_id=repository_id,
        trace_id=trace_id,
        api_key_id=api_key_id,
        metadata=metadata,
        created_at=datetime.now(UTC),
        current_username=current_username,
        llm_cost=llm_cost,
    )


async def get_or_create_optimization_event(
    trace_id: str,
    event_type: str = "no-pr",
    user_id: str | None = None,
    repository_id: str | None = None,
    repo_owner: str | None = None,
    repo_name: str | None = None,
    api_key_id: int | None = None,
    metadata: dict[str, Any] | None = None,
    current_username: str | None = None,
    llm_cost: float | None = None,
) -> tuple[OptimizationEvents, bool]:
    """Get existing optimization event or create a new one.

    Uses aupdate_or_create to handle race conditions atomically.

    Returns:
        A tuple of (event, created) where created is True if a new event was created.

    """
    if repository_id is None and repo_owner and repo_name:
        repository = await get_repository(repo_owner, repo_name)
        repository_id = str(repository.id) if repository else None

    defaults: dict[str, Any] = {}
    if repository_id is not None:
        defaults["repository_id"] = repository_id
    if current_username is not None:
        defaults["current_username"] = current_username

    event, created = await OptimizationEvents.objects.aupdate_or_create(
        trace_id=trace_id,
        defaults=defaults,
        create_defaults={
            "id": str(uuid4()),
            "event_type": event_type,
            "user_id": user_id,
            "repository_id": repository_id,
            "api_key_id": api_key_id,
            "metadata": metadata,
            "created_at": datetime.now(UTC),
            "current_username": current_username,
            "llm_cost": llm_cost,
        },
    )

    if not created and llm_cost is not None:
        await _increment_llm_cost(trace_id, llm_cost)

    return event, created


async def update_optimization_features_review(
    trace_id: str,
    review_quality: str | None = None,
    review_explanation: str | None = None,
    calling_fn_details: str | None = None,
) -> None:
    normalized = normalize_trace_id(trace_id)
    if normalized is None:
        return
    trace_uuid = UUID(normalized)
    # This avoids the race condition and is more performant.
    await OptimizationFeatures.objects.filter(trace_id=trace_uuid).aupdate(
        review_quality=review_quality, review_explanation=review_explanation, calling_fn_details=calling_fn_details
    )


async def update_optimization_cost(trace_id: str, cost: float, user_id: str | None = None) -> None:
    """Atomically increment llm_cost for the given trace_id.

    Creates a minimal event if one doesn't exist, then updates the cost.
    Uses aupdate_or_create to handle race conditions atomically.
    """
    _, created = await OptimizationEvents.objects.aupdate_or_create(
        trace_id=trace_id,
        defaults={},
        create_defaults={
            "id": str(uuid4()),
            "event_type": "no-pr",
            "user_id": user_id,
            "created_at": datetime.now(UTC),
            "llm_cost": cost,
        },
    )
    if not created:
        await _increment_llm_cost(trace_id, cost)
