import json
import time
import urllib.request

import oci


MODEL_ID = "openai.gpt-4o-mini"
MAX_EVENTS = 100


def get_instance_compartment() -> str:
    request = urllib.request.Request(
        "http://169.254.169.254/opc/v2/instance/",
        headers={"Authorization": "Bearer Oracle"},
    )

    with urllib.request.urlopen(
        request,
        timeout=5,
    ) as response:
        metadata = json.load(response)

    return metadata["compartmentId"]


def main() -> None:
    signer = (
        oci.auth.signers
        .InstancePrincipalsSecurityTokenSigner()
    )

    client = (
        oci.generative_ai_inference
        .GenerativeAiInferenceClient(
            config={"region": signer.region},
            signer=signer,
            timeout=(10, 60),
        )
    )

    models = oci.generative_ai_inference.models

    chat_request = models.GenericChatRequest(
        messages=[
            models.UserMessage(
                content=[
                    models.TextContent(
                        text=(
                            "Explique em português brasileiro, "
                            "em três frases curtas, como funciona "
                            "um assistente de voz."
                        )
                    )
                ]
            )
        ],
        is_stream=True,
        temperature=0.3,
        max_completion_tokens=160,
    )

    details = models.ChatDetails(
        compartment_id=get_instance_compartment(),
        serving_mode=models.OnDemandServingMode(
            model_id=MODEL_ID,
        ),
        chat_request=chat_request,
    )

    print("=" * 55)
    print("OCI GENERATIVE AI STREAMING TEST")
    print("=" * 55)
    print("Model:", MODEL_ID)
    print("Region:", signer.region)
    print()

    started = time.perf_counter()
    first_token_at = None
    event_count = 0

    response = None

    try:
        response = client.chat(details)

        print("HTTP status:", response.status)
        print("Response data type:", type(response.data).__name__)

        # OCI streaming responses expose an iterable event stream.
        stream = getattr(response.data, "events", None)

        if stream is None:
            raise RuntimeError(
                "OCI response does not expose an event stream."
            )

        for event in stream():
            event_count += 1

            elapsed = time.perf_counter() - started
            raw_data = getattr(event, "data", None)

            if event_count <= 5 or raw_data == "[DONE]":
                print(
                    f"[{elapsed:.3f}s] "
                    f"Event {event_count}: "
                    f"{str(raw_data)[:1000]}"
                )

            # This measures the first received event, not
            # necessarily the first generated text token.
            if first_token_at is None:
                first_token_at = elapsed

            if event_count >= MAX_EVENTS:
                print("Maximum event count reached.")
                break

    except oci.exceptions.ServiceError as exc:
        print("OCI SERVICE ERROR")
        print("HTTP status:", exc.status)
        print("Code:", exc.code)
        print("Message:", exc.message)
        raise

    finally:
        if response is not None:
            close = getattr(response.data, "close", None)

            if callable(close):
                close()

    total = time.perf_counter() - started

    print()
    print("=" * 55)
    print("RESULT")
    print("=" * 55)
    print("Events:", event_count)
    print("First event:", first_token_at)
    print("Total time:", round(total, 3), "seconds")


if __name__ == "__main__":
    main()