from __future__ import annotations

from typing import Any

from Mama.models import ConversationState, ExecutionContext, QueryResult
from Mama.ports import RuntimePorts

try:
    from langchain_core.output_parsers import StrOutputParser  # type: ignore
except Exception:  # pragma: no cover
    from langchain.schema import StrOutputParser  # type: ignore

try:
    from langchain_core.prompts import PromptTemplate  # type: ignore
except Exception:  # pragma: no cover
    from langchain.prompts import PromptTemplate  # type: ignore


def _history_to_text(messages: list[dict[str, Any]], max_messages: int) -> str:
    recent = messages[-max_messages:]
    lines: list[str] = []
    for message in recent:
        role = str(message.get("role", "user"))
        content = str(message.get("content", ""))
        lines.append(f"{role}: {content}")
    return "\n".join(lines)


def _append_message(state: ConversationState, role: str, content: str) -> ConversationState:
    messages = [*state.messages, {"role": role, "content": content}]
    if len(messages) > state.max_messages:
        messages = messages[-state.max_messages:]
    return ConversationState(messages=messages, max_messages=state.max_messages)


def query(
    ctx: ExecutionContext,
    question: str,
    runtime: RuntimePorts,
    *,
    prompt: str,
) -> QueryResult:
    text = (question or "").strip()
    if not text:
        raise ValueError("question cannot be empty")
    if not (prompt or "").strip():
        raise ValueError("prompt cannot be empty")

    embeddings = runtime.embeddings_provider.get_embeddings(ctx)
    store = runtime.vector_store_factory.get_store(ctx, embeddings)
    docs = store.similarity_search(
        text,
        k=ctx.retrieval.k,
        search_type=ctx.retrieval.search_type,
        filters=ctx.retrieval.filters,
    )

    context_text = "\n\n".join(doc.page_content for doc in docs)
    history_text = _history_to_text(ctx.conversation.messages, ctx.conversation.max_messages)

    llm = runtime.llm_provider.build_llm(ctx)
    prompt_template = PromptTemplate.from_template(
        "{prompt}\n\n"
        "Context:\n{context}\n"
        "Question: {question}\n"
        "Chat History:\n{chat_history}\n"
    )

    chain = prompt_template | llm | StrOutputParser()
    answer = chain.invoke(
        {
            "prompt": str(prompt).strip(),
            "context": context_text,
            "question": text,
            "chat_history": history_text,
        }
    )

    updated = _append_message(ctx.conversation, "user", text)
    updated = _append_message(updated, "assistant", str(answer))

    out_docs = []
    for doc in docs:
        out_docs.append(
            {
                "page_content": doc.page_content,
                "source": doc.metadata.get("source", ""),
                "metadata": dict(doc.metadata),
            }
        )

    return QueryResult(
        answer=str(answer),
        documents=out_docs,
        updated_conversation_state=updated,
        usage={"retrieved_documents": len(out_docs), "vector_store": store.identifier()},
    )
