"""Base chain builder — generic LCEL construction + safe invocation."""
from __future__ import annotations

import logging
from typing import Any

from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableSerializable
from langchain_core.runnables.history import RunnableWithMessageHistory
from pydantic import BaseModel

from claim_validator.chains.memory.manager import MemoryManager, MemoryType
from claim_validator.chains.provider_factory import build_chat_model

logger = logging.getLogger(__name__)


def build_chain(
    prompt: ChatPromptTemplate,
    output_schema: type[BaseModel],
    config: dict[str, Any],
    memory_manager: MemoryManager | None = None,
    memory_type: MemoryType = MemoryType.WITHIN_CLAIM,
    input_key: str = "claim_data",
) -> RunnableSerializable:
    """Build an LCEL chain: prompt | model | parser, optionally with memory."""
    model = build_chat_model(config)
    parser = PydanticOutputParser(pydantic_object=output_schema)
    chain: RunnableSerializable = prompt | model | parser

    if memory_manager is None:
        return chain

    return RunnableWithMessageHistory(
        chain,
        lambda session_id: memory_manager.get_history(session_id, memory_type),
        input_messages_key=input_key,
        history_messages_key="history",
    )


def invoke_chain_safe(
    chain: RunnableSerializable,
    input_data: dict[str, Any],
    config: dict[str, Any] | None = None,
    fallback_output: BaseModel | None = None,
) -> BaseModel | None:
    """Invoke a chain with 3-tier graceful degradation.

    1. Try chain (PydanticOutputParser)
    2. On OutputParserException, retry once
    3. On second failure, return fallback_output (or None)
    """
    from langchain_core.exceptions import OutputParserException

    try:
        return chain.invoke(input_data, config=config)
    except OutputParserException:
        logger.warning("Chain output parse failed, attempting retry", exc_info=True)
        try:
            return chain.invoke(input_data, config=config)
        except Exception:
            logger.warning("Chain retry also failed, using fallback", exc_info=True)
            if fallback_output is not None:
                return fallback_output
            return None
    except Exception:
        logger.warning("Chain invocation failed", exc_info=True)
        if fallback_output is not None:
            return fallback_output
        return None
