import os
import argparse
import asyncio
from typing import Union, Annotated
from pydantic import BaseModel, Field
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 langchain_core.output_parsers import PydanticOutputParser
from dotenv import load_dotenv

# Load API key from .env
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Pydantic models
class HttpOkEvent(BaseModel):
    kind: Literal["ok"] = Field(description="Event kind: ok")
    status: Literal[200] = Field(description="HTTP status code")
    path: str = Field(description="Request path")
    duration_ms: int = Field(description="Duration in milliseconds")

class HttpErrorEvent(BaseModel):
    kind: Literal["error"] = Field(description="Event kind: error")
    status: int = Field(description="HTTP status code")
    path: str = Field(description="Request path")
    error_message: str = Field(description="Error message")

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

# Structured output parser
parser = PydanticOutputParser(pydantic_object=ApiEvent)

# LLM configuration
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
)

# Backend for deepagents
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# System prompt for the agent
system_prompt = (
    "You are a log parser. "
    "Parse the given log line and output JSON that matches the following schema. "
    f"{parser.get_format_instructions()} "
    "Return only the JSON."
)

# Create the agent
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt=system_prompt,
)

async def parse_line(line: str) -> ApiEvent:
    """Parse a single log line using the agent and return a typed event."""
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=line)]},
        {"configurable": {"thread_id": "parser"}},
    )
    content = result["messages"][-1].content
    try:
        event = parser.parse(content)
    except Exception as e:
        raise ValueError(f"Failed to parse line: {line!r}. Error: {e}") from e
    return event

def print_table(events):
    """Print a simple table of events."""
    header = f"{'kind':<6} | {'path':<20} | {'status':<6} | {'detail':<30}"
    print(header)
    print("-" * len(header))
    for ev in events:
        if ev.kind == "ok":
            detail = f"duration {ev.duration_ms}ms"
        else:
            detail = ev.error_message
        print(f"{ev.kind:<6} | {ev.path:<20} | {ev.status:<6} | {detail:<30}")

async def main():
    parser_cli = argparse.ArgumentParser(description="Parse raw log lines into typed events.")
    parser_cli.add_argument(
        "--log",
        type=str,
        help="Raw log string (multiple lines). If omitted, a sample log is used.",
    )
    parser_cli.add_argument(
        "--file",
        type=str,
        help="Path to a file containing raw log lines.",
    )
    args = parser_cli.parse_args()

    if args.file:
        with open(args.file, "r", encoding="utf-8") as f:
            raw_log = f.read()
    elif args.log:
        raw_log = args.log
    else:
        raw_log = (
            "GET /api/users 200 123ms\n"
            "POST /api/login 404 Not Found\n"
            "GET /api/data 200 45ms\n"
            "DELETE /api/item/42 500 Internal Server Error\n"
            "PUT /api/update 200 78ms"
        )

    lines = [line.strip() for line in raw_log.splitlines() if line.strip()]
    events = []
    for line in lines:
        try:
            event = await parse_line(line)
            events.append(event)
            print(event.model_dump())
        except Exception as e:
            print(f"Error parsing line: {line!r}. {e}")

    print("\nParsed events table:")
    print_table(events)

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