# Hello ticket: start a session, send one message, poll for the reply, end the session.
# Python 3.8+ with "requests". Run: WETEL_API_KEY=sk_... WETEL_AGENT_ID=7 python hello_ticket.py
import os
import time

import requests

API_URL = os.environ.get("WETEL_API_URL", "https://api.wetel.dev/graphql")
API_KEY = os.environ["WETEL_API_KEY"]
AGENT_ID = int(os.environ.get("WETEL_AGENT_ID", "1"))
POLL_INTERVAL_SECONDS = 2
TIMEOUT_SECONDS = 60

START = """
mutation StartSession($input: SdkStartInput!) {
  sdkStart(input: $input) {
    sessionId
    avatarToken
    agentId
    graphqlEndpoint
  }
}
"""
SEND = """
mutation SendMessage($input: SdkSendMessageInput!) {
  sdkSendMessage(input: $input)
}
"""
POLL = """
query PollSession($externalCustomerId: String!, $first: Int) {
  sessionsByExternalCustomerId(externalCustomerId: $externalCustomerId, first: $first) {
    totalCount
    edges {
      node {
        id
        status
        messages {
          id
          role
          text
          turnIndex
          createdAt
        }
      }
    }
  }
}
"""
END = """
mutation EndSession($input: SdkEndSessionInput!) {
  sdkEndSession(input: $input) {
    sessionId
    durationSeconds
    evaluation {
      score
      summary
    }
  }
}
"""


def gql(query, variables, auth):
    response = requests.post(
        API_URL,
        json={"query": query, "variables": variables},
        headers={"x-huat-platform": "customer", **auth},
        timeout=30,
    )
    body = response.json()
    if body.get("errors"):
        error = body["errors"][0]
        raise RuntimeError(f"{error.get('extensions', {}).get('code')}: {error['message']}")
    return body["data"]


def reply_after_last_user_message(messages):
    """The reply is every MODEL message after the last USER message."""
    roles = [m["role"] for m in messages]
    last_user = len(roles) - 1 - roles[::-1].index("USER") if "USER" in roles else -1
    replies = [m["text"] for m in messages[last_user + 1 :] if m["role"] == "MODEL"]
    return "\n".join(replies) if replies else None


def main():
    agent_id = AGENT_ID
    client_id = f"ticket-{int(time.time() * 1000)}"  # one clientId per conversation
    text = "My laptop cannot connect to the office VPN since this morning."
    with_key = {"X-Api-Key": API_KEY}

    # 1. Start a session with the API key.
    started = gql(START, {
        "input": {
            "agentId": agent_id,
            "clientId": client_id,
        },
    }, with_key)["sdkStart"]
    session_id, avatar_token = started["sessionId"], started["avatarToken"]
    with_token = {"Authorization": f"Bearer {avatar_token}"}

    # 2. Send the message with the session token. It returns true at once.
    gql(SEND, {
        "input": {
            "sessionId": session_id,
            "text": text,
        },
    }, with_token)

    # 3. Poll the session's messages until the reply arrives.
    reply = None
    deadline = time.time() + TIMEOUT_SECONDS
    while reply is None and time.time() < deadline:
        time.sleep(POLL_INTERVAL_SECONDS)
        polled = gql(POLL, {
        "externalCustomerId": client_id,
        "first": 1,
    }, with_key)
        nodes = [edge["node"] for edge in polled["sessionsByExternalCustomerId"]["edges"]]
        session = next((n for n in nodes if n["id"] == session_id), None)
        reply = reply_after_last_user_message(session["messages"] if session else [])
    print(reply or "No reply within the timeout.")

    # 4. End the session with the session token.
    gql(END, {
        "input": {
            "sessionId": session_id,
        },
    }, with_token)


if __name__ == "__main__":
    main()