"""
# main.py – пример решения задания «Structured output – Union событий API»
# Использует deepagents, langchain-openai, pydantic v2 и стандартный argparse.
# Запуск: python main.py [--log TEXT]
# Если --log не указан, используется встроенный пример.
"""
import os
import argparse
import asyncio
from typing import Annotated, Literal, Union

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser

# ---------------------------------------------------------------------------
# 1. Модели Pydantic
# ---------------------------------------------------------------------------
class HttpOkEvent(BaseModel):
    kind: Literal["ok"] = Field("ok", description="Тип события – успешный запрос")
    status: Literal[200] = Field(200, description="HTTP‑код 200")
    path: str = Field(..., description="Запрошенный путь")
    duration_ms: int = Field(..., description="Время выполнения в миллисекундах")

class HttpErrorEvent(BaseModel):
    kind: Literal["error"] = Field("error", description="Тип события – ошибка")
    status: int = Field(..., description="HTTP‑код ошибки (4xx/5xx)")
    path: str = Field(..., description="Запрошенный путь")
    error_message: str = Field(..., description="Текст ошибки")

ApiEvent = Annotated[Union[HttpOkEvent, HttpErrorEvent], Field(discriminator="kind")]

# ---------------------------------------------------------------------------
# 2. LLM и агент
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# Парсер Pydantic
parser = PydanticOutputParser(pydantic_object=ApiEvent)

# ---------------------------------------------------------------------------
# 3. Инструмент – парсинг одного блока лога
# ---------------------------------------------------------------------------
@tool
def parse_log_block(block: str) -> str:
    """Парсит один блок лога и возвращает JSON‑строку, соответствующую ApiEvent."""
    # Формируем запрос к LLM
    prompt = (
        """
        Извлеките из следующего блока лога одно событие API и представьте его в JSON‑формате, соответствующем одной из двух схем:

        1. HttpOkEvent: {"kind": "ok", "status": 200, "path": "/api/...", "duration_ms": 123}
        2. HttpErrorEvent: {"kind": "error", "status": 404, "path": "/api/...", "error_message": "Not Found"}

        Ваша задача – не использовать regex, а просто сформировать корректный JSON, который можно распарсить Pydantic.
        """
        + "\n---\n"
        + block
        + "\n---"
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    json_text = response.content.strip()
    # Убедимся, что ответ начинается с { и заканчивается }
    if not (json_text.startswith("{") and json_text.endswith("}")):
        raise ValueError("LLM не вернул корректный JSON")
    return json_text

# ---------------------------------------------------------------------------
# 4. Агент – обрабатывает весь текст лога, разбивает на блоки и собирает список событий
# ---------------------------------------------------------------------------
agent = create_deep_agent(
    model=llm,
    tools=[parse_log_block],
    backend=backend,
    system_prompt="You are a log‑parser agent. Parse each log block into a structured event.",
)

async def process_log(text: str):
    # Разбиваем на блоки по строке "---" (или по пустой строке)
    blocks = [b.strip() for b in text.split("---") if b.strip()]
    events = []
    for block in blocks:
        # Запускаем агент для одного блока
        result = await agent.ainvoke(
            {"messages": [HumanMessage(content=block)]},
            {"configurable": {"thread_id": "log-session"}},
        )
        json_str = result["messages"][-1].content
        event = parser.parse(json_str)
        events.append(event)
    return events

# ---------------------------------------------------------------------------
# 5. CLI
# ---------------------------------------------------------------------------
DEFAULT_LOG = """
GET /api/users 200 123ms
---
POST /api/login 404 Not Found
---
GET /api/orders 500 Internal Server Error
---
PUT /api/users/42 200 45ms
"""

def main():
    parser_cli = argparse.ArgumentParser(description="Parse raw API logs into structured events.")
    parser_cli.add_argument("--log", type=str, help="Путь к файлу с логом или сам лог в кавычках")
    args = parser_cli.parse_args()

    if args.log:
        if os.path.isfile(args.log):
            with open(args.log, "r", encoding="utf-8") as f:
                raw_log = f.read()
        else:
            raw_log = args.log
    else:
        raw_log = DEFAULT_LOG

    events = asyncio.run(process_log(raw_log))

    # Выводим каждый объект в виде model_dump
    for ev in events:
        print(ev.model_dump())

    # Таблица
    print("\nТаблица: kind | path | status | ...")
    print("----------------------------------------")
    for ev in events:
        if ev.kind == "ok":
            print(f"{ev.kind:<4} | {ev.path:<20} | {ev.status:<5} | duration={ev.duration_ms}ms")
        else:
            print(f"{ev.kind:<4} | {ev.path:<20} | {ev.status:<5} | error='{ev.error_message}'")

if __name__ == "__main__":
    main()