import os
import asyncio
from typing import Annotated, Literal, Union

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

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

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

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

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

# ---------- Parser ----------
parser = PydanticOutputParser(pydantic_object=ApiEvent)

# ---------- Tool ----------
@tool
def parse_log_block(block: str) -> str:
    """Parse a single log block and return a JSON‑serialisable string of ApiEvent."""
    # The LLM will produce a JSON object that matches one of the Pydantic models.
    prompt = (
        "You are given a single log line.\n"
        "Return a JSON object that matches one of the following schemas:\n"
        "1. {\n  \"kind\": \"ok\", \"status\": 200, \"path\": \\"/api\\", \"duration_ms\": 123\n}\n"
        "2. {\n  \"kind\": \"error\", \"status\": 404, \"path\": \\"/api\\", \"error_message\": \"Not found\"\n}\n"
        "Do not add any extra keys.\n"
        f"Log line: {block}\n"
        "Answer in JSON only."
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    try:
        event = parser.parse(response.content)
        return event.model_dump_json()
    except Exception as e:
        return f"{{\"error\": \"{str(e)}\"}}"

# ---------- Agent ----------
agent = create_deep_agent(
    model=llm,
    tools=[parse_log_block],
    backend=backend,
    system_prompt="You are a log parsing assistant.",
)

# ---------- CLI logic ----------
DEFAULT_LOG = """
GET /api/users 200 123ms
POST /api/orders 404 Not Found
GET /api/products 200 45ms
"""

async def main():
    user_input = os.getenv("LOG_TEXT") or DEFAULT_LOG.strip()
    # Split into blocks by newlines, ignore empty
    blocks = [b for b in user_input.splitlines() if b.strip()]
    events = []
    for block in blocks:
        result = await agent.ainvoke(
            {"messages": [HumanMessage(content=f"Parse this block: {block}")], "tools": [parse_log_block]},
            {"configurable": {"thread_id": "log-session"}},
        )
        # The tool output is a JSON string; parse it
        try:
            event = parser.parse(result["messages"][-1].content)
            events.append(event)
        except Exception:
            continue
    # Print structured output
    for e in events:
        print(e.model_dump())
    # Table
    print("\nKind | Path | Status | Details")
    for e in events:
        if e.kind == "ok":
            print(f"ok | {e.path} | {e.status} | duration {e.duration_ms}ms")
        else:
            print(f"error | {e.path} | {e.status} | {e.error_message}")

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