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

from pydantic import BaseModel, Field, Literal
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

# ----------------- 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 (4xx/5xx)")
    path: str = Field(..., description="Request path")
    error_message: str = Field(..., description="Error message")

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

# ----------------- LLM -----------------
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 -----------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ----------------- Tools -----------------
@tool
def parse_block(block: str) -> ApiEvent:
    """Parse a single log block into a typed ApiEvent using LLM structured output."""
    # Use the LLM to produce structured output
    result = llm.with_structured_output(ApiEvent).invoke({"messages": [HumanMessage(content=block)]})
    # The LLM returns an ApiEvent instance directly
    return result

@tool
def parse_log(log_text: str) -> str:
    """Parse the entire log text and return a formatted table of events."""
    # Split into blocks (empty lines or '---' separators)
    raw_blocks = [b.strip() for b in log_text.split("\n") if b.strip()]
    events: List[ApiEvent] = []
    for block in raw_blocks:
        try:
            event = parse_block(block)
            events.append(event)
        except Exception as e:
            # If parsing fails, skip the block but keep a note
            events.append(
                HttpErrorEvent(
                    kind="error",
                    status=0,
                    path="<unknown>",
                    error_message=f"Failed to parse block: {e}",
                )
            )
    # Build table
    header = f"{'kind':<6} | {'path':<20} | {'status':<6} | {'detail':<30}"
    lines = [header, "-" * len(header)]
    for ev in events:
        if isinstance(ev, HttpOkEvent):
            detail = f"duration {ev.duration_ms}ms"
        else:
            detail = ev.error_message
        line = f"{ev.kind:<6} | {ev.path:<20} | {ev.status:<6} | {detail:<30}"
        lines.append(line)
    return "\n".join(lines)

# ----------------- Agent -----------------
agent = create_deep_agent(
    model=llm,
    tools=[parse_block, parse_log],
    backend=backend,
    system_prompt="You are a log parsing agent. Use the provided tools to parse the log and return a formatted table.",
)

# ----------------- CLI -----------------
async def run_agent(log_text: str):
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=log_text)]},
        {"configurable": {"thread_id": "log-session"}},
    )
    # The tool output is in the last message
    print(result["messages"][-1].content)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Parse raw log into structured events.")
    parser.add_argument("--file", type=str, help="Path to a log file. If omitted, uses default sample.")
    args = parser.parse_args()

    if args.file:
        with open(args.file, "r", encoding="utf-8") as f:
            log = f.read()
    else:
        # Default sample log
        log = """
2023-10-01 12:00:00 INFO /api/users 200 123ms
2023-10-01 12:00:01 ERROR /api/orders 404 Not Found
2023-10-01 12:00:02 INFO /api/products 200 98ms
2023-10-01 12:00:03 ERROR /api/payments 500 Internal Server Error
"""
    asyncio.run(run_agent(log))