app/services/wolfram_reference_resolver.py

from __future__ import annotations

from dataclasses import dataclass

from app.clients.wolfram import (
    WolframAlphaClient,
    WolframClientError,
    WolframQueryResult,
    WolframTextPod,
)
from app.schemas.cross_validation import EngineStatus


@dataclass(frozen=True)
class WolframReferenceResolution:
    query: str
    status: EngineStatus
    reference_result: str | None
    reason: str


_PREFERRED_POD_TITLES = {
    "result",
    "solution",
    "solutions",
    "root",
    "roots",
    "domain",
    "derivative",
    "limit",
    "exact result",
    "alternate forms",
    "decimal approximation",
}


def _clean_wolfram_plaintext(plaintext: str) -> str | None:
    cleaned = plaintext.strip()

    if not cleaned:
        return None

    # Wolfram retourne parfois plusieurs lignes. La première ligne utile suffit
    # pour une référence courte dans le MVP.
    useful_lines = [line.strip() for line in cleaned.splitlines() if line.strip()]
    if not useful_lines:
        return None

    return useful_lines[0]


def _extract_reference_result_from_pods(pods: list[WolframTextPod]) -> str | None:
    for pod in pods:
        title = pod.title.strip().lower()
        plaintext = _clean_wolfram_plaintext(pod.plaintext)

        if plaintext and title in _PREFERRED_POD_TITLES:
            return plaintext

    for pod in pods:
        plaintext = _clean_wolfram_plaintext(pod.plaintext)
        if plaintext:
            return plaintext

    return None


def resolve_wolfram_reference_from_result(
    wolfram_result: WolframQueryResult,
) -> WolframReferenceResolution:
    reference_result = _extract_reference_result_from_pods(wolfram_result.pods)

    if wolfram_result.success and reference_result:
        return WolframReferenceResolution(
            query=wolfram_result.query,
            status=EngineStatus.OK,
            reference_result=reference_result,
            reason="wolfram_reference_result_extracted",
        )

    if wolfram_result.error:
        return WolframReferenceResolution(
            query=wolfram_result.query,
            status=EngineStatus.FAILED,
            reference_result=None,
            reason="wolfram_query_error",
        )

    return WolframReferenceResolution(
        query=wolfram_result.query,
        status=EngineStatus.FAILED,
        reference_result=None,
        reason="wolfram_reference_result_missing",
    )


async def resolve_wolfram_reference(query: str) -> WolframReferenceResolution:
    cleaned_query = query.strip()

    if not cleaned_query:
        return WolframReferenceResolution(
            query="",
            status=EngineStatus.FAILED,
            reference_result=None,
            reason="wolfram_query_empty",
        )

    try:
        wolfram_result = await WolframAlphaClient().query(cleaned_query)
    except WolframClientError as exc:
        return WolframReferenceResolution(
            query=cleaned_query,
            status=EngineStatus.FAILED,
            reference_result=None,
            reason=str(exc),
        )

    return resolve_wolfram_reference_from_result(wolfram_result)