"""Compose taste, intent, grounding, and reranking into streamed events."""

import asyncio
import time
from collections.abc import AsyncGenerator, Iterator
from contextlib import aclosing

import structlog

from app.adapters.spotify.errors import SpotifyError
from app.config import Settings
from app.domain.models import (
    CompressedTasteProfile,
    ConversationTurn,
    Intent,
    PreviousRecommendation,
    TasteProfile,
    Track,
)
from app.domain.profile import compress_taste_profile
from app.observability.timing import increment_cache_hits
from app.pipeline.event import (
    PipelineDoneEvent,
    PipelineErrorEvent,
    PipelineEvent,
    PipelineMetadataEvent,
    PipelineTrackEvent,
    PipelineWarningEvent,
)
from app.pipeline.grounding import Grounder
from app.ports.protocols import (
    CatalogQuotaExhaustedError,
    MusicCatalog,
    Recommender,
    RecommenderOutputError,
)

RERANK_FALLBACK_CODE = "rerank_fallback"
RERANK_FALLBACK_MESSAGE = "Ranking output was invalid, so grounded results are shown instead."
RERANK_FALLBACK_JUSTIFICATION = "Selected as a grounded match for your request."


class _TrackSelection:
    """Rank tracks from one grounded pool, enforcing bound and uniqueness."""

    def __init__(self, pool: tuple[Track, ...], limit: int) -> None:
        """Bind the only tracks that may ever be selected."""
        self.pool = pool
        self.limit = limit
        self.selected: list[Track] = []
        self._tracks_by_id = {track.id: track for track in pool}
        self._selected_ids: set[str] = set()

    @property
    def is_full(self) -> bool:
        """Return whether the selection reached its limit."""
        return len(self.selected) >= self.limit

    @property
    def remaining_count(self) -> int:
        """Return how many further selections are allowed."""
        return self.limit - len(self.selected)

    def describe_selected_ids(self) -> str:
        """Render the ids already emitted, for a correction instruction."""
        return ", ".join(sorted(self._selected_ids)) or "none"

    def select(self, track_id: str, justification: str) -> PipelineTrackEvent:
        """Accept one recommender selection or reject it as invalid output."""
        if self.is_full:
            raise RecommenderOutputError("Rerank returned too many track ids")
        track = self._tracks_by_id.get(track_id)
        if track is None:
            raise RecommenderOutputError("Rerank selected an out-of-pool track id")
        if track_id in self._selected_ids:
            raise RecommenderOutputError("Rerank selected a duplicate track id")
        return self._emit(track, justification)

    def fill_from_pool(self, justification: str) -> Iterator[PipelineTrackEvent]:
        """Complete the selection in pool order after a failed rerank."""
        for track in self.pool:
            if self.is_full:
                return
            if track.id not in self._selected_ids:
                yield self._emit(track, justification)

    def _emit(self, track: Track, justification: str) -> PipelineTrackEvent:
        self._selected_ids.add(track.id)
        self.selected.append(track)
        return PipelineTrackEvent(
            rank=len(self.selected),
            track=track,
            justification=justification,
        )


class RecommendationPipeline:
    """Orchestrate the code-defined recommendation stages."""

    def __init__(
        self,
        recommender: Recommender,
        settings: Settings,
        grounder: Grounder | None = None,
    ) -> None:
        """Create process-local caches around the provided service ports."""
        self.recommender = recommender
        self.settings = settings
        self.grounder = grounder or Grounder(settings)
        self.taste_cache = _TasteProfileCache(settings)
        self.last_pools: dict[str, tuple[Track, ...]] = {}

    async def stream(
        self,
        session_id: str,
        request_id: str,
        catalog: MusicCatalog,
        query: str,
        history: tuple[ConversationTurn, ...],
        previous_recommendations: tuple[PreviousRecommendation, ...],
    ) -> AsyncGenerator[PipelineEvent]:
        """Yield ordered events for one recommendation request."""
        started_at = time.monotonic()
        deadline_at = started_at + self.settings.request_deadline_seconds
        try:
            taste = await self.taste_cache.get(session_id, catalog)
            intent = await self.recommender.create_intent(
                query,
                history,
                previous_recommendations,
                taste.text,
                self.settings.candidate_count,
            )
        except RecommenderOutputError:
            yield PipelineErrorEvent(
                code="intent_failed",
                message="Recommendation intent could not be generated from the model response.",
            )
            return
        except CatalogQuotaExhaustedError:
            yield PipelineErrorEvent(
                code="quota_exhausted",
                message=(
                    "Spotify request quota was exhausted before recommendations could be prepared."
                ),
            )
            return
        except SpotifyError:
            yield PipelineErrorEvent(
                code="spotify_unavailable",
                message="Spotify was unavailable while preparing recommendations.",
            )
            return
        yield PipelineMetadataEvent(
            request_id=request_id,
            intent_summary=intent.intent_summary,
            candidate_count=len(intent.candidates),
        )

        pool = await self._grounded_pool(session_id, catalog, intent, taste, deadline_at)
        if not pool:
            yield PipelineErrorEvent(
                code="no_grounded_results",
                message="None of the proposed tracks could be verified on Spotify.",
            )
            return
        if len(pool) < self.settings.grounding_floor:
            # Fewer verified tracks than promised is still an answer; an
            # empty error in its place would hide real results.
            yield PipelineWarningEvent(
                code="partial_results",
                message="Fewer tracks than usual could be verified; showing what held up.",
            )

        selection = _TrackSelection(pool, self.settings.rerank_count)
        async for event in self._ranked_events(intent, taste.text, history, selection):
            yield event

        _log_completion(selection, taste)
        yield PipelineDoneEvent(
            track_count=len(selection.selected),
            total_ms=round((time.monotonic() - started_at) * 1000),
        )

    async def _grounded_pool(
        self,
        session_id: str,
        catalog: MusicCatalog,
        intent: Intent,
        taste: CompressedTasteProfile,
        deadline_at: float,
    ) -> tuple[Track, ...]:
        """Reuse the session's pool on refinement, otherwise ground anew."""
        if intent.is_refinement:
            cached_pool = self.last_pools.get(session_id)
            if cached_pool:
                return cached_pool
        result = await self.grounder.ground(
            catalog,
            intent.candidates,
            taste.known_track_ids,
            intent.familiarity,
            self.settings.rerank_count + self.settings.rerank_pool_buffer,
            deadline_at,
        )
        if result.tracks:
            self.last_pools[session_id] = result.tracks
        return result.tracks

    async def _ranked_events(
        self,
        intent: Intent,
        taste_summary: str,
        history: tuple[ConversationTurn, ...],
        selection: _TrackSelection,
    ) -> AsyncGenerator[PipelineTrackEvent | PipelineWarningEvent]:
        """Stream the rerank with one corrected retry, then fall back."""
        try:
            async for event in self._rerank_with_one_retry(
                intent, taste_summary, history, selection
            ):
                yield event
            return
        except RecommenderOutputError as error:
            _log_rerank_failure(attempt=2, error=error)

        yield PipelineWarningEvent(code=RERANK_FALLBACK_CODE, message=RERANK_FALLBACK_MESSAGE)
        for event in selection.fill_from_pool(RERANK_FALLBACK_JUSTIFICATION):
            yield event

    async def _rerank_with_one_retry(
        self,
        intent: Intent,
        taste_summary: str,
        history: tuple[ConversationTurn, ...],
        selection: _TrackSelection,
    ) -> AsyncGenerator[PipelineTrackEvent]:
        """Rerank once; on invalid output, retry once with a correction."""
        try:
            async for event in self._rerank_once(intent, taste_summary, history, selection, None):
                yield event
            return
        except RecommenderOutputError as error:
            _log_rerank_failure(attempt=1, error=error)
            if selection.is_full:
                return
            correction = (
                f"Validation failed: {error}."
                f" Already emitted track ids: {selection.describe_selected_ids()}."
            )

        async for event in self._rerank_once(intent, taste_summary, history, selection, correction):
            yield event

    async def _rerank_once(
        self,
        intent: Intent,
        taste_summary: str,
        history: tuple[ConversationTurn, ...],
        selection: _TrackSelection,
        correction: str | None,
    ) -> AsyncGenerator[PipelineTrackEvent]:
        stream = self.recommender.stream_rerank(
            intent,
            selection.pool,
            taste_summary,
            history,
            selection.remaining_count,
            correction,
        )
        async with aclosing(stream) as selections:
            async for item in selections:
                yield selection.select(item.track_id, item.justification)


def _log_rerank_failure(attempt: int, error: RecommenderOutputError) -> None:
    """Log one invalid rerank attempt with its validation reason."""
    structlog.get_logger().warning("rerank_attempt_failed", attempt=attempt, error=str(error))


def _log_completion(selection: _TrackSelection, taste: CompressedTasteProfile) -> None:
    """Log how many recommendations were served and how many are new."""
    new_track_count = sum(track.id not in taste.known_track_ids for track in selection.selected)
    structlog.get_logger().info(
        "recommendations_complete",
        track_count=len(selection.selected),
        new_track_count=new_track_count,
    )


class _TasteProfileCache:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self._entries: dict[str, tuple[float, CompressedTasteProfile]] = {}
        self._locks: dict[str, asyncio.Lock] = {}

    async def get(self, session_id: str, catalog: MusicCatalog) -> CompressedTasteProfile:
        cached = self._fresh_entry(session_id)
        if cached is not None:
            increment_cache_hits()
            return cached
        lock = self._locks.setdefault(session_id, asyncio.Lock())
        async with lock:
            cached = self._fresh_entry(session_id)
            if cached is not None:
                increment_cache_hits()
                return cached
            compressed = await self._fetch(catalog)
            self._entries[session_id] = (time.monotonic(), compressed)
            return compressed

    def _fresh_entry(self, session_id: str) -> CompressedTasteProfile | None:
        entry = self._entries.get(session_id)
        if entry is None:
            return None
        created_at, profile = entry
        if time.monotonic() - created_at >= self.settings.taste_profile_ttl_seconds:
            del self._entries[session_id]
            return None
        return profile

    async def _fetch(self, catalog: MusicCatalog) -> CompressedTasteProfile:
        try:
            async with asyncio.TaskGroup() as task_group:
                short_artists_task = task_group.create_task(
                    catalog.fetch_top_artists("short_term", self.settings.top_items_limit)
                )
                long_artists_task = task_group.create_task(
                    catalog.fetch_top_artists("long_term", self.settings.top_items_limit)
                )
                short_tracks_task = task_group.create_task(
                    catalog.fetch_top_tracks("short_term", self.settings.top_items_limit)
                )
                long_tracks_task = task_group.create_task(
                    catalog.fetch_top_tracks("long_term", self.settings.top_items_limit)
                )
                saved_tracks_task = task_group.create_task(
                    catalog.fetch_saved_tracks(self.settings.saved_tracks_limit)
                )
        except ExceptionGroup as errors:
            raise errors.exceptions[0] from None
        return compress_taste_profile(
            TasteProfile(
                short_term_artists=tuple(short_artists_task.result()),
                long_term_artists=tuple(long_artists_task.result()),
                short_term_tracks=tuple(short_tracks_task.result()),
                long_term_tracks=tuple(long_tracks_task.result()),
                saved_tracks=tuple(saved_tracks_task.result()),
            )
        )