import os
import argparse
from typing import Annotated, Literal, Union, List

from pydantic import BaseModel, Field, ValidationError
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

# --------------------------------------------------------------------------- #
# 1. Модели событий
# --------------------------------------------------------------------------- #

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

class HttpErrorEvent(BaseModel):
    kind: Literal["error"] = Field(..., 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 + Pydantic
# --------------------------------------------------------------------------- #

parser = PydanticOutputParser(pydantic_object=ApiEvent)

prompt_template = PromptTemplate(
    input_variables=["log_line"],
    template=(
        "Найди в следующем логе событие API и выведи его в формате JSON, "
        "соответствующем одной из следующих схем:\n"
        "1. {ok_schema}\n"
        "2. {error_schema}\n"
        "Укажи поле \"kind\" как \"ok\" или \"error\".\n"
        "Лог: {log_line}\n"
        "Ответ в чистом JSON без комментариев."
    ),
)

# --------------------------------------------------------------------------- #
# 3. Функция обработки одной строки
# --------------------------------------------------------------------------- #

def parse_line(line: str, llm: ChatOpenAI) -> ApiEvent:
    ok_schema = HttpOkEvent.model_json_schema()
    error_schema = HttpErrorEvent.model_json_schema()
    prompt = prompt_template.format(
        log_line=line,
        ok_schema=ok_schema,
        error_schema=error_schema,
    )
    response = llm.invoke(prompt)
    try:
        event = parser.parse(response.content)
    except ValidationError as e:
        raise ValueError(f"LLM не смог распарсить строку: {line}\n{e}") from e
    return event

# --------------------------------------------------------------------------- #
# 4. Обработка всего лога
# --------------------------------------------------------------------------- #

def parse_log(text: str, llm: ChatOpenAI) -> List[ApiEvent]:
    events: List[ApiEvent] = []
    # Разделяем по строкам, игнорируем пустые
    for line in text.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        event = parse_line(line, llm)
        events.append(event)
    return events

# --------------------------------------------------------------------------- #
# 5. CLI
# --------------------------------------------------------------------------- #

def main():
    parser_cli = argparse.ArgumentParser(description="Parse API logs into structured events.")
    parser_cli.add_argument(
        "--log",
        type=str,
        default=None,
        help="Путь к файлу с логом. Если не задан, используется пример.",
    )
    args = parser_cli.parse_args()

    if args.log:
        with open(args.log, "r", encoding="utf-8") as f:
            log_text = f.read()
    else:
        log_text = """\
2023-10-01 12:00:01 INFO /api/users 200 123ms
2023-10-01 12:00:02 ERROR /api/orders 404 Not Found
2023-10-01 12:00:03 WARN /api/payments 500 Internal Server Error
"""

    # Инициализируем LLM
    openai_api_key = os.getenv("OPENAI_API_KEY")
    if not openai_api_key:
        raise RuntimeError("Переменная окружения OPENAI_API_KEY не найдена")
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=openai_api_key)

    events = parse_log(log_text, llm)

    # Выводим каждый объект
    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} | {ev.path} | {ev.status} |")
        else:
            print(f"| {ev.kind} | {ev.path} | {ev.status} |")

if __name__ == "__main__":
    main()