import os
import argparse
from typing import Annotated, Union
from pydantic import BaseModel, Field, ValidationError, Literal
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import PydanticOutputParser
from dotenv import load_dotenv

# Load environment variables for API key
load_dotenv()

# Define Pydantic models for log events
class HttpOkEvent(BaseModel):
    kind: Literal["ok"] = "ok"
    status: int = Field(..., description="HTTP status code")
    path: str

class HttpErrorEvent(BaseModel):
    kind: Literal["error"] = "error"
    status: int = Field(..., description="HTTP status code")
    path: str
    error: str

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

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

# LLM instance
llm = ChatOpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL"),
    model=os.getenv("OPENAI_MODEL", "gpt-3.5-turbo"),
)


def parse_line(line: str) -> ApiEvent:
    """Parse a single log line using LLM and structured output parser."""
    prompt = (
        f"You are a log parser. {parser.get_format_instructions()}\n"
        f"Log line: {line}\n"
    )
    try:
        raw_output = llm.invoke(prompt)
        if hasattr(raw_output, "content"):
            raw_text = raw_output.content
        else:
            raw_text = str(raw_output)
        event = parser.parse(raw_text)
        return event
    except Exception as e:
        # Fallback: simple regex parsing
        # This fallback is only for safety and not part of primary implementation.
        # It attempts to extract status, path, and error if present.
        # If parsing fails, raise the original exception.
        import re
        error_pattern = r"(?P<status>\d{3})\s+(?P<path>\S+)(?:\s+error:\s+(?P<error>.+))?"
        match = re.search(error_pattern, line)
        if match:
            groups = match.groupdict()
            status = int(groups["status"])
            path = groups["path"]
            if status >= 400:
                return HttpErrorEvent(kind="error", status=status, path=path, error=groups.get("error", ""))
            else:
                return HttpOkEvent(kind="ok", status=status, path=path)
        else:
            raise e


def summarize_events(events: list[ApiEvent]) -> None:
    """Print a summary table of events."""
    header = f"{'Kind':<6} {'Status':<6} {'Path':<20} {'Error':<30}"
    print(header)
    print('-' * len(header))
    for e in events:
        if isinstance(e, HttpOkEvent):
            print(f"{e.kind:<6} {e.status:<6} {e.path:<20} {'':<30}")
        elif isinstance(e, HttpErrorEvent):
            print(f"{e.kind:<6} {e.status:<6} {e.path:<20} {e.error:<30}")


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

    if args.input:
        if os.path.isfile(args.input):
            with open(args.input, "r", encoding="utf-8") as f:
                data = f.read()
        else:
            data = args.input
    else:
        # sample log
        data = (
            "200 /api/v1/users\n"
            "404 /api/v1/items error: Not Found\n"
            "500 /api/v1/orders error: Internal Server Error"
        )

    # Split into lines or blocks
    lines = [line for line in data.splitlines() if line.strip()]
    events: list[ApiEvent] = []
    for line in lines:
        try:
            event = parse_line(line)
            events.append(event)
        except Exception as exc:
            print(f"Failed to parse line: {line!r}. Error: {exc}")

    if events:
        summarize_events(events)
    else:
        print("No events parsed.")


if __name__ == "__main__":
    main()