import logging
import time
from dataclasses import dataclass, field
from datetime import datetime
from types import SimpleNamespace

from config import calculate_cost
from rag.rag_usage import RAGWithUsage

logger = logging.getLogger(__name__)


@dataclass
class LLMCallRecord:
    model: str
    prompt: str
    instructions: str
    answer: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    response_time: float
    cost: float
    timestamp: datetime = field(default_factory=datetime.now)


class RAGWithMetrics(RAGWithUsage):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.last_call: LLMCallRecord = None

    def llm(self, prompt):
        start_time = time.time()
        # Call parent's llm() which handles _call_llm + usage tracking
        answer = super().llm(prompt)
        response_time = time.time() - start_time

        # last_usage is set by super().llm(); normalise missing usage
        usage = self.last_usage
        if usage is None:
            usage = SimpleNamespace(
                prompt_tokens=0, completion_tokens=0, total_tokens=0
            )

        cost = calculate_cost(self.model, usage)
        call_record = LLMCallRecord(
            model=self.model,
            prompt=prompt,
            instructions=self.instructions,
            answer=answer,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            total_tokens=usage.total_tokens,
            response_time=response_time,
            cost=cost,
        )
        logger.info("LLM call: %s", call_record)
        self.last_call = call_record
        return answer