"""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.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 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()
        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,
        )
        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)
        if len(pool) < self.settings.grounding_floor:
            yield PipelineErrorEvent(
                code="insufficient_grounding",
                message="Not enough requested tracks could be verified safely.",
            )
            return

        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,
    ) -> 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,
        )
        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 one rerank, retry once on invalid output, then fall back."""
        correction: str | None = None
        for _ in range(2):
            if selection.is_full:
                return
            try:
                async for event in self._rerank_once(
                    intent, taste_summary, history, selection, correction
                ):
                    yield event
                return
            except RecommenderOutputError as error:
                correction = (
                    f"Validation failed: {error}."
                    f" Already emitted track ids: {selection.describe_selected_ids()}."
                )

        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_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_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:
        short_artists, long_artists, short_tracks, long_tracks, saved_tracks = await asyncio.gather(
            catalog.fetch_top_artists("short_term", self.settings.top_items_limit),
            catalog.fetch_top_artists("long_term", self.settings.top_items_limit),
            catalog.fetch_top_tracks("short_term", self.settings.top_items_limit),
            catalog.fetch_top_tracks("long_term", self.settings.top_items_limit),
            catalog.fetch_saved_tracks(self.settings.saved_tracks_limit),
        )
        return compress_taste_profile(
            TasteProfile(
                short_term_artists=tuple(short_artists),
                long_term_artists=tuple(long_artists),
                short_term_tracks=tuple(short_tracks),
                long_term_tracks=tuple(long_tracks),
                saved_tracks=tuple(saved_tracks),
            )
        )