"""Stream-режим AI-агента: потоковый вывод токенов и событий updates."""
from __future__ import annotations

import os

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI

load_dotenv()

BROJS_INFERENCE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1"
DEFAULT_MODEL = "openai/gpt-oss-20b:free"
STEP_SEPARATOR = "\n --- --- --- \n"

DEFAULT_QUESTION = (
    "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
)


def _api_key() -> str:
    return (
        os.getenv("OPENAI_API_KEY")
        or os.getenv("JOURNAL_MCP_PAT")
        or os.getenv("JOURNAL_TOKEN")
        or ""
    )


def _base_url() -> str:
    if os.getenv("OPENAI_BASE_URL"):
        return os.environ["OPENAI_BASE_URL"]
    if os.getenv("OPENAI_API_KEY"):
        return os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
    return BROJS_INFERENCE_URL


def _model() -> str:
    return os.getenv("OPENAI_MODEL") or os.getenv("OPENROUTER_MODEL") or DEFAULT_MODEL


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=_model(),
        base_url=_base_url(),
        api_key=_api_key(),
        temperature=0.7,
    )


def _extract_table(text: str) -> str:
    lines = [line for line in text.splitlines() if "|" in line]
    if lines:
        return "\n".join(lines)
    return text.strip()


def _build_price_subagent(llm: ChatOpenAI):
    return create_agent(
        model=llm,
        system_prompt=(
            "Ты аналитик цен на продукты питания. "
            "По названию продукта и городу оцени реалистичную цену в рублях. "
            "Ответь ТОЛЬКО одной строкой markdown-таблицы:\n"
            "| Продукт | Цена (руб.) | Магазин |"
        ),
    )


def make_get_price_tool(llm: ChatOpenAI):
    price_subagent = _build_price_subagent(llm)

    @tool
    def get_price(product: str, city: str) -> str:
        """Возвращает примерную цену продукта в указанном городе.

        Args:
            product: название продукта (молоко, хлеб, яблоки и т.д.)
            city: город покупателя
        """
        prompt = (
            f"Город: {city}. Продукт: {product}. "
            "Верни одну строку таблицы с ценой и магазином."
        )
        result = price_subagent.invoke(
            {"messages": [{"role": "human", "content": prompt}]}
        )
        last = result["messages"][-1]
        content = getattr(last, "content", str(last))
        return _extract_table(content)

    return get_price


def build_agent(llm: ChatOpenAI | None = None):
    """Главный агент покупок — точка входа для автопроверки."""
    llm = llm or build_llm()
    get_price = make_get_price_tool(llm)
    return create_agent(
        model=llm,
        tools=[get_price],
        system_prompt="Ты помощник по планированию покупок",
    )


def format_message(message) -> str:
    """Текст сообщения или вызов инструмента."""
    content = getattr(message, "content", None)
    if content:
        return str(content)
    tool_calls = getattr(message, "tool_calls", None) or []
    if tool_calls:
        call = tool_calls[0]
        name = call.get("name") if isinstance(call, dict) else getattr(call, "name", "")
        args = call.get("args") if isinstance(call, dict) else getattr(call, "args", {})
        return f"{name}({args})"
    return str(message)


def run_shopping_assistant_stream(
    agent=None,
    question: str = DEFAULT_QUESTION,
) -> None:
    """Потоковый запуск агента (messages + updates)."""
    agent = agent or build_agent()

    stream = agent.stream(
        {"messages": [{"role": "human", "content": question}]},
        stream_mode=["messages", "updates"],
    )

    step = 1

    def format_chunk_message(chunk_data: tuple) -> None:
        nonlocal step
        message, meta = chunk_data
        graph_step = meta.get("langgraph_step", step)
        if graph_step != step:
            step = graph_step
            print(STEP_SEPARATOR, end="")
        if message.content:
            print(message.content, end="", flush=True)

    for chunk in stream:
        chunk_type, chunk_data = chunk

        if chunk_type == "messages":
            format_chunk_message(chunk_data)

        if chunk_type == "updates":
            model_update = chunk_data.get("model")
            if model_update:
                last_message = model_update["messages"][-1]
                formatted = format_message(last_message)
                if formatted.strip():
                    print(formatted)

    print()


def main() -> None:
    run_shopping_assistant_stream()


if __name__ == "__main__":
    main()