import asyncio
import os
import uuid
from typing import Any, Dict, Optional

import atexit
import click
import structlog
import time
from anthropic.types import MessageParam as AnthropicMessageParam
from google.genai import types as gdm_types

from openreward.models import (
    LogMessageEvent,
    RolloutConfig,
    RolloutStartedEvent,
    SendLoopConfig,
    ShutdownEvent,
)
from .background import start_background_worker
from .serializers.ant import serialize_anthropic_message
from .serializers.base import UploadType, base_to_normalized
from .serializers.gdm import serialize_gdm_message
from .serializers.models import NormalizedEvent
from .serializers.oai_completions import (
    OpenAIChatMessage,
    openai_completions_to_normalized,
)
from .serializers.oai_responses import serialize_openai_response

logger = structlog.get_logger("openreward.rollouts.rollout")
OPENREWARD_DISABLE_ROLLOUT_LOGS = bool(os.getenv("OPENREWARD_DISABLE_ROLLOUT_LOGS", False))

class RolloutAPI:
    def __init__(self,
        send_loop_config: SendLoopConfig,
        shutdown_timeout: float = 10.0,
        web_base_url: str = "https://openreward.ai",
    ):
        self.send_loop_config = send_loop_config
        self.shutdown_timeout = shutdown_timeout
        self.web_base_url = web_base_url

        self._loop, self._in, self._thread = start_background_worker(send_loop_config)
        self._closed = False

        atexit.register(self._atexit_handler)

    def create(
        self,
        run_name: str,
        rollout_name: Optional[str] = None,
        environment: Optional[str] = None,
        split: Optional[str] = None,
        metadata: Optional[dict] = None,
        task_spec: Optional[Dict[str, Any]] = None
    ) -> "Rollout":
        config = RolloutConfig(
            run_name=run_name,
            rollout_name=rollout_name,
            environment=environment,
            split=split,
            metadata=metadata,
            task_spec=task_spec
        )
        return Rollout(config=config, _in=self._in, _loop=self._loop, web_base_url=self.web_base_url)

    def close(self):
        if self._closed:
            return
        self._closed = True
        if self._thread.is_alive():
            self._loop.call_soon_threadsafe(self._in.put_nowait, ShutdownEvent())
            self._thread.join(timeout=self.shutdown_timeout)
            if self._thread.is_alive():
                print("openreward: shutdown timed out, some rollout data may not have been uploaded")

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def _atexit_handler(self):
        # Best effort if user forgot
        try:
            self.close()
        except Exception:
            pass

class Rollout:
    def __init__(self, config: RolloutConfig, _in: asyncio.Queue, _loop: asyncio.AbstractEventLoop, web_base_url: str = "https://openreward.ai"):
        self.config = config
        self.event_id = str(uuid.uuid4())
        self._in = _in
        self._loop = _loop
        self._loop.call_soon_threadsafe(self._in.put_nowait, RolloutStartedEvent.from_config(
            event_id=self.event_id,
            timestamp=int(time.time() * 1000),
            config=config
        ))
        self.logged_messages = 0

        url = f"{web_base_url.rstrip('/')}/rollout/{self.event_id}"
        if not OPENREWARD_DISABLE_ROLLOUT_LOGS:
            prefix = click.style("openreward:", bold=True, fg="blue")
            recording = click.style("Recording rollout", fg=(247, 230, 204))
            styled_url = click.style(url, fg="blue", underline=True)
            click.echo(f"{prefix} \U0001f680 {recording} {styled_url}", err=True)

    def _log_message(self, normalized_event: NormalizedEvent, reward: Optional[float] = None, is_finished: Optional[bool] = False, metadata: Optional[dict] = None):

        if is_finished is None:
            is_finished = False

        self._loop.call_soon_threadsafe(self._in.put_nowait, LogMessageEvent(
            eventId = str(uuid.uuid4()),
            timestamp=int(time.time() * 1000),
            index=self.logged_messages,
            rolloutEventId=self.event_id,
            type=normalized_event.type,
            content=normalized_event.content,
            contentReference=normalized_event.content_reference,
            summary=normalized_event.summary,
            name=normalized_event.name,
            callId=normalized_event.call_id,
            reward=reward,
            is_finished=is_finished,
            metadata=metadata or {},
        ))
        self.logged_messages += 1

    def log(
        self,
        message: UploadType,
        reward: Optional[float] = None,
        is_finished: Optional[bool] = False,
        metadata: Optional[dict] = None,
    ):
        normalized_event = base_to_normalized([message])[0]
        self._log_message(normalized_event, reward, is_finished, metadata)

    def log_openai_completions(
        self,
        message: OpenAIChatMessage,
        reward: Optional[float] = None,
        is_finished: Optional[bool] = False,
        metadata: Optional[dict] = None,
    ):
        normalized_event = openai_completions_to_normalized([message])[0]
        self._log_message(normalized_event, reward, is_finished, metadata)

    def log_openai_response(
        self,
        message,  # Accept both Response and ResponseInputItemParam
        reward: Optional[float] = None,
        is_finished: Optional[bool] = False,
        metadata: Optional[dict] = None,
    ):
        # Handle Response objects by iterating over output
        messages_to_log = []
        if hasattr(message, 'output'):
            messages_to_log = list(message.output)
        else:
            messages_to_log = [message]

        for msg in messages_to_log:
            event = serialize_openai_response(msg)
            if event:  # Skip None (shouldn't happen with fixes, but safe)
                self._log_message(event, reward, is_finished, metadata)

    def log_anthropic_message(
        self,
        message: AnthropicMessageParam,
        reward: Optional[float] = None,
        is_finished: Optional[bool] = False,
        metadata: Optional[dict] = None,
    ):
        for event in serialize_anthropic_message(message):
            self._log_message(event, reward, is_finished, metadata)

    def log_gdm_message(
        self,
        message: gdm_types.Content,
        reward: Optional[float] = None,
        is_finished: Optional[bool] = False,
        metadata: Optional[dict] = None,
    ):
        for event in serialize_gdm_message(message):
            self._log_message(event, reward, is_finished, metadata)
