#!/usr/bin/env python
"""Structured log parser using DeepAgents and Pydantic v2.

This script demonstrates how to parse a raw log containing HTTP
responses into a list of typed events.  It uses the DeepAgents
framework to create a lightweight agent that delegates the parsing
job to an OpenRouter LLM via structured output.  The result is a
list of Pydantic models that can be printed or displayed in a table.

Usage:
    python main.py                # uses built‑in demo log
    python main.py "<your log>"   # parse custom log passed as a single string
"""

import os
import sys
import asyncio
from typing import Annotated, Literal, Union, List
from pathlib import Path

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
from tabulate import tabulate

# ---------------------------------------------------------------------------
# 1. Pydantic models – Union with discriminator
# ---------------------------------------------------------------------------

class HttpOkEvent(BaseModel):
    kind: Literal["ok"] = Field("ok", description="Event kind – OK")
    status: Literal[200] = Field(200, description="HTTP status code")
    path: str = Field(..., description="Request path")
    duration_ms: int = Field(..., description="Response time in milliseconds")

class HttpErrorEvent(BaseModel):
    kind: Literal["error"] = Field("error", description="Event kind – Error")
    status: int = Field(..., description="HTTP status code (4xx/5xx)")
    path: str = Field(..., description="Request path")
    error_message: str = Field(..., description="Error description")

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

# ---------------------------------------------------------------------------
# 2. LLM and Agent setup
# ---------------------------------------------------------------------------

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(),
])

# The agent will simply forward the prompt to the LLM and return the
# structured output.  No additional tools are required for this task.
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a log‑parsing assistant.  Return a structured
    representation of the event using the provided Pydantic models.",
)

# ---------------------------------------------------------------------------
# 3. Parsing helper
# ---------------------------------------------------------------------------

parser = PydanticOutputParser(pydantic_object=ApiEvent)

async def parse_block(block: str) -> ApiEvent:
    """Ask the LLM to parse a single log block into an ApiEvent.

    The LLM is instructed to output only the JSON that matches the
    Pydantic schema.  The parser then validates and returns the model.
    """
    prompt = (
        "Parse the following log entry and return a JSON object that matches "
        "one of the following schemas: HttpOkEvent or HttpErrorEvent.  "
        "Do not include any additional keys or text.
        """
        f"Log entry:\n{block.strip()}"
    )
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": "parser"}},
    )
    content = response["messages"][-1].content
    return parser.parse(content)

# ---------------------------------------------------------------------------
# 4. Main logic
# ---------------------------------------------------------------------------

DEFAULT_LOG = """"""
DEFAULT_LOG += "GET /api/users 200 123ms\n"
DEFAULT_LOG += "POST /api/login 404 Not Found\n"
DEFAULT_LOG += "GET /api/data 500 Internal Server Error\n"
DEFAULT_LOG += "PUT /api/update 200 98ms\n"
DEFAULT_LOG += "DELETE /api/remove 403 Forbidden\n"

async def main():
    raw_log = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_LOG
    # Split into blocks – simple newline split but ignore empty lines
    blocks = [b for b in raw_log.strip().split("\n") if b]
    events: List[ApiEvent] = []
    for block in blocks:
        try:
            event = await parse_block(block)
            events.append(event)
        except Exception as e:
            print(f"Failed to parse block: {block}\nError: {e}")
    # Output each event as JSON
    for ev in events:
        print(ev.model_dump_json(indent=2))
    # Pretty table
    table = [[
        ev.kind,
        ev.path,
        ev.status,
        getattr(ev, "duration_ms", "-"),
        getattr(ev, "error_message", "-"),
    ] for ev in events]
    headers = ["kind", "path", "status", "duration_ms", "error_message"]
    print("\nParsed events table:\n")
    print(tabulate(table, headers=headers, tablefmt="github"))

if __name__ == "__main__":
    asyncio.run(main())