from typing import TYPE_CHECKING, Dict, List

import click
from rich.progress import Progress

from ....libs.players.params import PlayerParams
from ....libs.providers.anime.params import EpisodeStreamsParams
from ..session import Context, session
from ..state import ControlFlow, ProviderState, State

if TYPE_CHECKING:
    from ....cli.utils.utils import (
        filter_by_quality,  # You may need to create this helper
    )
    from ....libs.providers.anime.types import Server


def _filter_by_quality(links, quality):
    # Simplified version of your filter_by_quality for brevity
    for link in links:
        if str(link.quality) == quality:
            return link
    return links[0] if links else None


@session.menu
def servers(ctx: Context, state: State) -> State | ControlFlow:
    """
    Fetches and displays available streaming servers for a chosen episode,
    then launches the media player and transitions to post-playback controls.
    """
    provider_anime = state.provider.anime
    episode_number = state.provider.episode_number
    config = ctx.config
    provider = ctx.provider
    selector = ctx.selector

    if not provider_anime or not episode_number:
        click.echo("[bold red]Error: Anime or episode details are missing.[/bold red]")
        return ControlFlow.BACK

    # --- Fetch Server Streams ---
    with Progress(transient=True) as progress:
        progress.add_task(
            f"[cyan]Fetching servers for episode {episode_number}...", total=None
        )
        server_iterator = provider.episode_streams(
            EpisodeStreamsParams(
                anime_id=provider_anime.id,
                episode=episode_number,
                translation_type=config.stream.translation_type,
            )
        )
        # Consume the iterator to get a list of all servers
        all_servers: List[Server] = list(server_iterator) if server_iterator else []

    if not all_servers:
        click.echo(
            f"[bold yellow]No streaming servers found for this episode.[/bold yellow]"
        )
        return ControlFlow.BACK

    # --- Auto-Select or Prompt for Server ---
    server_map: Dict[str, Server] = {s.name: s for s in all_servers}
    selected_server: Server | None = None

    preferred_server = config.stream.server.lower()
    if preferred_server == "top":
        selected_server = all_servers[0]
        click.echo(f"[cyan]Auto-selecting top server:[/] {selected_server.name}")
    elif preferred_server in server_map:
        selected_server = server_map[preferred_server]
        click.echo(f"[cyan]Auto-selecting preferred server:[/] {selected_server.name}")
    else:
        choices = [*server_map.keys(), "Back"]
        chosen_name = selector.choose("Select Server", choices)
        if not chosen_name or chosen_name == "Back":
            return ControlFlow.BACK
        selected_server = server_map[chosen_name]

    if not selected_server:
        return ControlFlow.BACK

    # --- Select Stream Quality ---
    stream_link_obj = _filter_by_quality(selected_server.links, config.stream.quality)
    if not stream_link_obj:
        click.echo(
            f"[bold red]No stream of quality '{config.stream.quality}' found on server '{selected_server.name}'.[/bold red]"
        )
        return ControlFlow.CONTINUE

    # --- Launch Player ---
    final_title = f"{provider_anime.title} - Ep {episode_number}"
    click.echo(f"[bold green]Launching player for:[/] {final_title}")

    player_result = ctx.player.play(
        PlayerParams(
            url=stream_link_obj.link,
            title=final_title,
            subtitles=[sub.url for sub in selected_server.subtitles],
            headers=selected_server.headers,
            # start_time logic will be added in player_controls
        )
    )

    # --- Transition to Player Controls ---
    # We now have all the data for post-playback actions.
    return State(
        menu_name="PLAYER_CONTROLS",
        media_api=state.media_api,
        provider=state.provider.model_copy(
            update={
                "servers": all_servers,
                "selected_server": selected_server,
                "last_player_result": player_result,  # We should add this to ProviderState
            }
        ),
    )