import argparse
import os
from typing import Annotated, Literal, Union, List
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv

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

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

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

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

# ---------------------------------------------------------------------------
# 2. LLM and Parser setup
# ---------------------------------------------------------------------------

load_dotenv()
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, max_output_tokens=512)
parser = PydanticOutputParser(pydantic_object=ApiEvent)
chain = RunnablePassthrough() | llm | parser

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

def parse_log_line(line: str) -> ApiEvent:
    prompt = f"Parse the following log line into JSON:\n{line}\nThe JSON should match one of the following schemas:\n- ok event: {{\"kind\": \"ok\", status: 200, path: string, duration_ms: int}}\n- error event: {{\"kind\": \"error\", status: int, path: string, error_message: string}}\nReturn only the JSON."
    result = chain.invoke(prompt)
    return result

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

def main():
    parser_cli = argparse.ArgumentParser(description="Parse log events into typed objects.")
    parser_cli.add_argument("--log", type=str, help="Path to log file or raw log string.")
    args = parser_cli.parse_args()

    if args.log:
        if os.path.exists(args.log):
            with open(args.log, "r", encoding="utf-8") as f:
                raw = f.read()
        else:
            raw = args.log
    else:
        raw = """GET /api/users 200 123ms
POST /api/users 404 Not Found
GET /api/orders 500 Internal Server Error
"""

    lines = [l.strip() for l in raw.splitlines() if l.strip()]
    events: List[ApiEvent] = []
    for line in lines:
        try:
            event = parse_log_line(line)
            events.append(event)
        except Exception as e:
            print(f"Failed to parse line: {line}\nError: {e}")

    print("\nParsed Events:")
    for ev in events:
        print(ev.model_dump())

    print("\nTable:\")
    header = ["kind", "path", "status"]
    print("{:<6} {:<20} {:<6}".format(*header))
    for ev in events:
        kind = ev.kind
        path = getattr(ev, "path", "")
        status = getattr(ev, "status", "")
        print("{:<6} {:<20} {:<6}".format(kind, path, status))

if __name__ == "__main__":
    main()