"""Lineage recording for inference results.

Extracts persistence fields from InferenceResult so that actual endpoint,
model, and route lineage are recorded — fixing the hardcoded
``model_provider = 'ollama'`` pattern.

Requirements: 2.9, 13.6
"""
from __future__ import annotations

from services.shared.inference.models import InferenceResult, ModelLineage


def build_lineage_from_result(result: InferenceResult, trace_id: str = "") -> ModelLineage:
    """Extract a ModelLineage record from an InferenceResult.

    This replaces hardcoded ``model_provider = 'ollama'`` persistence
    by capturing the actual endpoint, deployment, model, protocol,
    structured mode, request ID, latency, and retries from the result.

    Args:
        result: The completed inference result.
        trace_id: Optional distributed trace ID for correlation.

    Returns:
        A ModelLineage with all required persistence fields.
    """
    return ModelLineage(
        endpoint_id=result.endpoint_id,
        deployment_id=result.deployment_id,
        model=result.model,
        protocol=result.protocol,
        structured_mode=result.structured_mode,
        request_id=result.request_id,
        latency_ms=result.latency_ms,
        retries=result.retries,
        trace_id=trace_id,
    )


def lineage_to_persistence_dict(lineage: ModelLineage) -> dict:
    """Convert a ModelLineage to a flat dict for database persistence.

    Returns fields suitable for inserting into agent_performance_log,
    document_intelligence, or similar tables.

    The ``model_provider`` field is derived from the protocol:
    - ``ollama_native`` → ``"ollama"``
    - ``openai_chat`` → ``"openai_compatible"`` (covers vLLM, OpenAI, etc.)
    - ``specialist_http`` → ``"specialist"``
    """
    protocol_to_provider = {
        "ollama_native": "ollama",
        "openai_chat": "openai_compatible",
        "specialist_http": "specialist",
    }

    return {
        "model_provider": protocol_to_provider.get(lineage.protocol, lineage.protocol),
        "model_name": lineage.model,
        "endpoint_id": str(lineage.endpoint_id) if lineage.endpoint_id else None,
        "deployment_id": str(lineage.deployment_id) if lineage.deployment_id else None,
        "protocol": lineage.protocol,
        "structured_mode": lineage.structured_mode,
        "request_id": lineage.request_id,
        "latency_ms": lineage.latency_ms,
        "retries": lineage.retries,
        "trace_id": lineage.trace_id,
    }