from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Any

from .models import LLMCall, OptimizationError

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam

    from aiservice.llm import LLMResponse


async def record_llm_call(
    trace_id: str,
    call_type: str,
    model_name: str,
    messages: list[ChatCompletionMessageParam],
    latency_ms: int,
    user_id: str | None = None,
    python_version: str | None = None,
    context: dict[str, Any] | None = None,
    result: LLMResponse | None = None,
    error: Exception | None = None,
    llm_cost: float | None = None,
) -> str:
    """Record LLM call with result. Returns llm_call_id."""
    llm_call_id = str(uuid.uuid4())
    system_prompt = next((m["content"] for m in messages if m["role"] == "system"), "")
    user_prompt = "\n\n".join(m["content"] for m in messages if m["role"] == "user" and m.get("content"))

    raw_response = None
    if result and hasattr(result.raw_response, "model_dump_json"):
        raw_response = result.raw_response.model_dump_json(indent=2)

    await LLMCall.objects.acreate(
        id=llm_call_id,
        trace_id=trace_id,
        call_type=call_type,
        model_name=model_name,
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        messages=messages,
        user_id=user_id,
        python_version=python_version,
        context=context,
        status="success" if result else "failed",
        raw_response=raw_response,
        prompt_tokens=result.usage.input_tokens if result else None,
        completion_tokens=result.usage.output_tokens if result else None,
        llm_cost=llm_cost,
        latency_ms=latency_ms,
        error_type=type(error).__name__ if error else None,
        error_message=str(error) if error else None,
    )
    return llm_call_id


async def record_error(
    trace_id: str,
    error_type: str,
    error_message: str,
    error_category: str = "llm_error",
    severity: str = "error",
    error_code: str | None = None,
    context: dict[str, Any] | None = None,
) -> str:
    """Record error. Returns error_id."""
    error_id = str(uuid.uuid4())
    await OptimizationError.objects.acreate(
        id=error_id,
        trace_id=trace_id,
        error_type=error_type,
        error_category=error_category,
        severity=severity,
        error_message=error_message,
        error_code=error_code,
        context=context,
    )
    return error_id
