"""Generate and synchronize avatars for MindRoom entities."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from functools import cache
from typing import TYPE_CHECKING, Literal

from google import genai
from google.genai import types
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.text import Text

from mindroom import constants
from mindroom.config.main import Config, load_config
from mindroom.constants import ROUTER_AGENT_NAME, resolve_avatar_path, workspace_avatar_path
from mindroom.credentials_sync import get_secret_from_env
from mindroom.error_handling import AvatarGenerationError, AvatarSyncError
from mindroom.logging_config import get_logger
from mindroom.matrix.avatar import room_has_avatar, set_room_avatar_from_file
from mindroom.matrix.identity import MatrixID
from mindroom.matrix.state import MatrixAccount, MatrixState, get_room_id, matrix_state_for_runtime
from mindroom.matrix.users import AgentMatrixUser, login_agent_user
from mindroom.matrix_identifiers import extract_server_name_from_homeserver
from mindroom.model_defaults import GOOGLE_AVATAR_IMAGE, GOOGLE_AVATAR_PROMPT

if TYPE_CHECKING:
    from pathlib import Path

    import nio


logger = get_logger(__name__)

_PROMPT_MODEL = GOOGLE_AVATAR_PROMPT
_IMAGE_MODEL = GOOGLE_AVATAR_IMAGE
_ROOT_SPACE_AVATAR_NAME = "root_space"
# Team prompts include per-member breakdowns; a low cap truncates them
# mid-sentence and the mangled prompt makes the image model answer with
# text instead of an image.
_PROMPT_MAX_OUTPUT_TOKENS = 400
# The image model occasionally returns no image for a valid prompt, so retry
# the complete prompt and image-generation request.
_MAX_IMAGE_ATTEMPTS = 3

_ROOM_PURPOSES = {
    "lobby": "Central meeting space, entrance and welcome area",
    "research": "Scientific investigation and data analysis",
    "docs": "Documentation and writing center",
    "ops": "Operations and system management",
    "automation": "Workflow automation and bot control",
    "analysis": "Data analysis and insights",
    "business": "Business strategy and planning",
    "communication": "Messages and team communication",
    "dev": "Software development and coding",
    "finance": "Financial analysis and trading",
    "help": "Support and assistance center",
    "home": "Personal home base and dashboard",
    "news": "News updates and current events",
    "productivity": "Task management and efficiency",
    "science": "Scientific research and experiments",
}

_AvatarEntityType = Literal["agents", "teams", "rooms", "spaces"]


@dataclass(frozen=True)
class _AvatarTeamMember:
    """A typed team member descriptor used for prompt generation."""

    name: str
    role: str


@dataclass(frozen=True)
class _AvatarTarget:
    """A typed avatar generation target."""

    entity_type: _AvatarEntityType
    entity_name: str
    role: str
    team_members: tuple[_AvatarTeamMember, ...] = ()


@cache
def _get_console() -> Console:
    """Create the shared Rich console used by avatar generation output."""
    return Console()


def _load_validated_config(runtime_paths: constants.RuntimePaths) -> Config:
    """Load and validate the active MindRoom configuration."""
    return load_config(runtime_paths, tolerate_plugin_load_errors=True)


def _get_avatar_path(
    entity_type: str,
    entity_name: str,
    runtime_paths: constants.RuntimePaths,
) -> Path:
    """Get the output path for an avatar file."""
    avatar_path = workspace_avatar_path(entity_type, entity_name, runtime_paths)
    avatar_path.parent.mkdir(parents=True, exist_ok=True)
    return avatar_path


def _managed_room_avatar_keys(config: Config) -> set[str]:
    """Return room keys that participate in managed avatar generation and sync."""
    return {room_name for room_name in config.get_all_configured_rooms() if not room_name.startswith(("!", "#"))}


def _managed_avatar_targets(config: Config) -> list[tuple[str, str]]:
    """Return every managed avatar target for the active config."""
    targets = [("agents", agent_name) for agent_name in config.agents]
    targets.append(("agents", "router"))
    targets.extend(("teams", team_name) for team_name in config.teams)
    targets.extend(("rooms", room_name) for room_name in _managed_room_avatar_keys(config))
    if config.matrix_space.enabled:
        targets.append(("spaces", _ROOT_SPACE_AVATAR_NAME))
    return targets


def _missing_avatar_targets(
    config: Config,
    runtime_paths: constants.RuntimePaths,
) -> set[tuple[str, str]]:
    """Return the managed avatar targets with no bundled or workspace avatar yet."""
    return {
        (entity_type, entity_name)
        for entity_type, entity_name in _managed_avatar_targets(config)
        if not resolve_avatar_path(entity_type, entity_name, runtime_paths).exists()
    }


async def _generate_prompt(
    client: genai.Client,
    target: _AvatarTarget,
    config: Config,
) -> str:
    """Generate an image prompt based on the entity's role using AI."""
    if target.entity_type in {"rooms", "spaces"}:
        system_prompt = config.get_prompt("AVATAR_ROOM_SYSTEM_PROMPT")
        user_prompt = f"Room name: {target.entity_name}\nPurpose: {target.role}"
    elif target.entity_type == "teams":
        system_prompt = config.get_prompt("AVATAR_TEAM_SYSTEM_PROMPT")
        user_prompt = f"Team name: {target.entity_name}\nTeam role: {target.role}"
        if target.team_members:
            members_info = "\n".join(f"- {member.name}: {member.role}" for member in target.team_members)
            user_prompt = f"{user_prompt}\nTeam members:\n{members_info}"
    else:
        system_prompt = config.get_prompt("AVATAR_AGENT_SYSTEM_PROMPT")
        user_prompt = f"Agent name: {target.entity_name}\nRole: {target.role}\nType: {target.entity_type}"

    response = await client.aio.models.generate_content(
        model=_PROMPT_MODEL,
        contents=user_prompt,
        config=types.GenerateContentConfig(
            system_instruction=system_prompt,
            max_output_tokens=_PROMPT_MAX_OUTPUT_TOKENS,
        ),
    )
    if not response.text:
        msg = f"Gemini returned no text prompt for {target.entity_type}/{target.entity_name}"
        raise ValueError(msg)

    visual_elements = response.text.strip()
    base_style = (
        config.get_prompt("AVATAR_ROOM_STYLE")
        if target.entity_type in {"rooms", "spaces"}
        else config.get_prompt("AVATAR_CHARACTER_STYLE")
    )
    final_prompt = f"{base_style}, {visual_elements}"

    _get_console().print(
        Panel(
            Text(final_prompt, style="cyan"),
            title=f"[bold yellow]{target.entity_type}/{target.entity_name}[/bold yellow]",
            border_style="green",
        ),
    )
    return final_prompt


def _no_image_diagnostic(response: types.GenerateContentResponse) -> str:
    """Summarize a no-image response for logs without dumping the payload."""
    finish_reasons = [
        str(candidate.finish_reason) for candidate in response.candidates or [] if candidate.finish_reason
    ]
    texts = [text for part in response.parts or [] if isinstance(text := getattr(part, "text", None), str)]
    snippet = " ".join(texts).strip()[:200]
    return f"finish_reasons={finish_reasons or None} text={snippet!r}"


def _extract_image_bytes(response: types.GenerateContentResponse) -> bytes | None:
    """Return the first generated image bytes from a Gemini response."""
    for part in response.parts or []:
        if part.inline_data and part.inline_data.data:
            return part.inline_data.data
    return None


async def _generate_avatar(
    client: genai.Client,
    target: _AvatarTarget,
    runtime_paths: constants.RuntimePaths,
    config: Config,
    *,
    force: bool = False,
) -> None:
    """Generate an avatar for a single entity if it does not exist."""
    avatar_path = _get_avatar_path(target.entity_type, target.entity_name, runtime_paths)
    if avatar_path.exists() and not force:
        _get_console().print(
            f"[green]✓[/green] Avatar already exists for [bold]{target.entity_type}/{target.entity_name}[/bold]",
        )
        return

    console = _get_console()
    console.print(f"\n[yellow]🎨 Generating avatar for {target.entity_type}/{target.entity_name}...[/yellow]")
    console.print(f"   [dim]Role: {target.role}[/dim]")
    if target.team_members:
        console.print(f"   [dim]Team members: {', '.join(member.name for member in target.team_members)}[/dim]")

    image_bytes: bytes | None = None
    for attempt in range(1, _MAX_IMAGE_ATTEMPTS + 1):
        prompt = await _generate_prompt(client, target, config)
        response = await client.aio.models.generate_content(
            model=_IMAGE_MODEL,
            contents=prompt,
            config=types.GenerateContentConfig(
                response_modalities=["IMAGE"],
                image_config=types.ImageConfig(
                    aspect_ratio="1:1",
                    image_size="1K",
                ),
            ),
        )
        image_bytes = _extract_image_bytes(response)
        if image_bytes:
            break
        logger.warning(
            "Image model returned no image for avatar target",
            entity_type=target.entity_type,
            entity_name=target.entity_name,
            attempt=attempt,
            max_attempts=_MAX_IMAGE_ATTEMPTS,
            response_diagnostic=_no_image_diagnostic(response),
        )
        if attempt < _MAX_IMAGE_ATTEMPTS:
            console.print(
                f"[yellow]⟳ No image returned for {target.entity_type}/{target.entity_name} "
                f"(attempt {attempt}/{_MAX_IMAGE_ATTEMPTS}); regenerating prompt[/yellow]",
            )

    if not image_bytes:
        msg = f"No image data found for {target.entity_type}/{target.entity_name} after {_MAX_IMAGE_ATTEMPTS} attempts"
        raise ValueError(msg)

    avatar_path.write_bytes(image_bytes)
    console.print(f"[green]✓ Generated avatar for {target.entity_type}/{target.entity_name}[/green]")


def _build_router_user(
    router_account: MatrixAccount,
    runtime_paths: constants.RuntimePaths,
) -> AgentMatrixUser:
    """Create the router user object from persisted Matrix state."""
    server_name = extract_server_name_from_homeserver(
        constants.runtime_matrix_homeserver(runtime_paths=runtime_paths),
        runtime_paths=runtime_paths,
    )
    return AgentMatrixUser(
        agent_name=ROUTER_AGENT_NAME,
        user_id=MatrixID.from_username(router_account.username, router_account.domain or server_name).full_id,
        display_name="Router",
        password=router_account.password,
        access_token=None,
    )


async def _sync_avatar_target(
    client: nio.AsyncClient,
    *,
    avatar_path: Path,
    room_id: str,
    label: str,
    force: bool = False,
) -> bool | None:
    """Apply one managed avatar target unless the room already has an avatar."""
    if not force and await room_has_avatar(client, room_id):
        _get_console().print(f"[dim]⊘ Skipped avatar for {label} (already set)[/dim]")
        return None

    if await set_room_avatar_from_file(client, room_id, avatar_path):
        _get_console().print(f"[green]✓ Set avatar for {label}[/green]")
        return True
    _get_console().print(f"[red]✗ Failed to set avatar for {label}[/red]")
    return False


async def _sync_configured_room_avatars(
    client: nio.AsyncClient,
    config: Config,
    runtime_paths: constants.RuntimePaths,
    *,
    force: bool = False,
) -> tuple[int, int, list[str]]:
    """Apply configured room avatars and return success/skip counts plus failed labels."""
    success_count = 0
    skip_count = 0
    failed_labels: list[str] = []
    for room_name in sorted(_managed_room_avatar_keys(config)):
        avatar_path = resolve_avatar_path("rooms", room_name, runtime_paths)
        if not avatar_path.exists():
            skip_count += 1
            continue

        room_id = get_room_id(room_name, runtime_paths)
        if not room_id:
            _get_console().print(f"[yellow]⚠ Room '{room_name}' not found in Matrix[/yellow]")
            continue

        label = f"room '{room_name}'"
        success = await _sync_avatar_target(
            client,
            avatar_path=avatar_path,
            room_id=room_id,
            label=label,
            force=force,
        )
        if success is True:
            success_count += 1
        elif success is None:
            skip_count += 1
        else:
            failed_labels.append(label)
    return success_count, skip_count, failed_labels


async def _sync_root_space_avatar(
    client: nio.AsyncClient,
    config: Config,
    state: MatrixState,
    runtime_paths: constants.RuntimePaths,
    *,
    force: bool = False,
) -> bool | None:
    """Apply the managed root-space avatar when both the asset and room exist."""
    if not config.matrix_space.enabled or not state.space_room_id:
        return None

    root_space_avatar_path = resolve_avatar_path(
        "spaces",
        _ROOT_SPACE_AVATAR_NAME,
        runtime_paths,
    )
    if not root_space_avatar_path.exists():
        return None

    return await _sync_avatar_target(
        client,
        avatar_path=root_space_avatar_path,
        room_id=state.space_room_id,
        label="root space",
        force=force,
    )


async def set_room_avatars_in_matrix(runtime_paths: constants.RuntimePaths, *, force: bool = False) -> None:
    """Set avatars for all rooms in Matrix."""
    console = _get_console()
    console.print("\n[bold cyan]Setting room avatars in Matrix...[/bold cyan]")

    state = matrix_state_for_runtime(runtime_paths)
    router_account = state.get_account(f"agent_{ROUTER_AGENT_NAME}")
    if not router_account:
        msg = "No router account found in Matrix state. Make sure mindroom has been started at least once."
        raise AvatarSyncError(msg)

    router_user = _build_router_user(router_account, runtime_paths)
    try:
        client = await login_agent_user(
            constants.runtime_matrix_homeserver(runtime_paths=runtime_paths),
            router_user,
            runtime_paths,
        )
    except ValueError as exc:
        msg = f"Failed to log in as router for avatar sync: {exc}"
        raise AvatarSyncError(msg) from exc
    console.print("[green]✓ Logged in to Matrix as router[/green]")

    config = _load_validated_config(runtime_paths)
    failed_labels: list[str] = []
    try:
        success_count, skip_count, failed_labels = await _sync_configured_room_avatars(
            client,
            config,
            runtime_paths,
            force=force,
        )
        root_space_success = await _sync_root_space_avatar(
            client,
            config,
            state,
            runtime_paths,
            force=force,
        )
        if root_space_success is True:
            success_count += 1
        elif root_space_success is False:
            failed_labels.append("root space")
    finally:
        await client.close()

    if success_count > 0:
        console.print(f"\n[green]✓ Set {success_count} room avatars[/green]")
    if skip_count > 0:
        console.print(f"[dim]⊘ Skipped {skip_count} rooms (no avatar file or avatar already set)[/dim]")
    if failed_labels:
        formatted_labels = ", ".join(failed_labels)
        msg = f"Failed to set avatars for: {formatted_labels}"
        console.print(f"[red]Error:[/red] {msg}")
        raise AvatarSyncError(msg)


def _build_avatar_generation_targets(
    config: Config,
    missing_targets: set[tuple[str, str]],
) -> list[_AvatarTarget]:
    """Build typed avatar targets for every missing managed avatar."""
    targets: list[_AvatarTarget] = []

    for agent_name, agent_config in config.agents.items():
        if ("agents", agent_name) in missing_targets:
            targets.append(
                _AvatarTarget(
                    entity_type="agents",
                    entity_name=agent_name,
                    role=agent_config.role or "AI assistant",
                ),
            )

    if ("agents", "router") in missing_targets:
        targets.append(
            _AvatarTarget(
                entity_type="agents",
                entity_name="router",
                role="Intelligent routing and agent or team selection",
            ),
        )

    for team_name, team_config in config.teams.items():
        if ("teams", team_name) in missing_targets:
            team_members = tuple(
                _AvatarTeamMember(
                    name=agent_name,
                    role=config.agents[agent_name].role or "Team member",
                )
                for agent_name in team_config.agents
                if agent_name in config.agents
            )
            targets.append(
                _AvatarTarget(
                    entity_type="teams",
                    entity_name=team_name,
                    role=team_config.role,
                    team_members=team_members,
                ),
            )

    targets.extend(
        _AvatarTarget(
            entity_type="rooms",
            entity_name=room_name,
            role=_ROOM_PURPOSES.get(room_name, f"Collaboration space for {room_name} activities"),
        )
        for room_name in _managed_room_avatar_keys(config)
        if ("rooms", room_name) in missing_targets
    )

    if ("spaces", _ROOT_SPACE_AVATAR_NAME) in missing_targets:
        targets.append(
            _AvatarTarget(
                entity_type="spaces",
                entity_name=_ROOT_SPACE_AVATAR_NAME,
                role=f"Workspace space named {config.matrix_space.name} that organizes all managed rooms",
            ),
        )

    return targets


def _print_avatar_generation_plan(missing_targets: set[tuple[str, str]]) -> None:
    """Print the number of missing avatars that will be generated."""
    space_count = int(("spaces", _ROOT_SPACE_AVATAR_NAME) in missing_targets)
    room_count = sum(1 for entity_type, _ in missing_targets if entity_type == "rooms")
    team_count = sum(1 for entity_type, _ in missing_targets if entity_type == "teams")
    agent_count = sum(1 for entity_type, _ in missing_targets if entity_type == "agents")
    _get_console().print(
        f"\n[bold cyan]🚀 Generating {agent_count} agents, {team_count} teams, {room_count} rooms, and {space_count} spaces...[/bold cyan]\n",
    )


def _remaining_missing_avatar_targets(
    missing_targets: set[tuple[str, str]],
    runtime_paths: constants.RuntimePaths,
) -> set[tuple[str, str]]:
    """Return targets that are still missing after a generation attempt."""
    return {
        (entity_type, entity_name)
        for entity_type, entity_name in missing_targets
        if not workspace_avatar_path(entity_type, entity_name, runtime_paths).exists()
    }


async def _generate_missing_avatars(
    config: Config,
    runtime_paths: constants.RuntimePaths,
    selected_targets: set[tuple[str, str]],
    *,
    force: bool = False,
) -> bool:
    """Generate every missing managed avatar and report whether startup may continue."""
    console = _get_console()
    if not selected_targets:
        console.print("\n[dim]⊘ All managed avatars already exist; skipping generation[/dim]")
        return True

    api_key = get_secret_from_env("GOOGLE_API_KEY", runtime_paths=runtime_paths)
    if not api_key:
        console.print("[red]Error: GOOGLE_API_KEY or GOOGLE_API_KEY_FILE environment variable not set[/red]")
        console.print("Please set it in your .env file, secrets mount, or environment")
        return False

    client = genai.Client(api_key=api_key)
    targets = _build_avatar_generation_targets(config, selected_targets)
    _print_avatar_generation_plan(selected_targets)

    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
        ) as progress:
            task_id = progress.add_task("Processing avatars...", total=None)
            results = await asyncio.gather(
                *(
                    _generate_avatar(
                        client,
                        target,
                        runtime_paths,
                        config,
                        force=force,
                    )
                    for target in targets
                ),
                return_exceptions=True,
            )
            progress.update(task_id, completed=True)
    finally:
        await client.aio.aclose()

    failed_targets: list[tuple[_AvatarTarget, Exception]] = []
    for target, result in zip(targets, results, strict=True):
        if isinstance(result, Exception):
            failed_targets.append((target, result))
            logger.error(
                "Avatar generation failed",
                entity_type=target.entity_type,
                entity_name=target.entity_name,
                error=repr(result),
                exc_info=(type(result), result, result.__traceback__),
            )
            console.print(f"[red]✗ Failed to generate {target.entity_type}/{target.entity_name}: {result}[/red]")

    remaining_targets = _remaining_missing_avatar_targets(selected_targets, runtime_paths)
    if failed_targets or remaining_targets:
        failed_target_keys = {(target.entity_type, target.entity_name) for target, _error in failed_targets}
        formatted_targets = ", ".join(
            f"{entity_type}/{entity_name}"
            for entity_type, entity_name in sorted(remaining_targets | failed_target_keys)
        )
        console.print(f"\n[red]✗ Avatar generation failed for: {formatted_targets}[/red]")
        return False

    console.print("\n[bold green]✨ Avatar generation complete![/bold green]")
    return True


async def run_avatar_generation(runtime_paths: constants.RuntimePaths, *, force: bool = False) -> None:
    """Generate missing managed avatars in the workspace."""
    config = _load_validated_config(runtime_paths)
    selected_targets = set(_managed_avatar_targets(config)) if force else _missing_avatar_targets(config, runtime_paths)

    if not await _generate_missing_avatars(config, runtime_paths, selected_targets, force=force):
        msg = "Avatar generation failed. See errors above."
        raise AvatarGenerationError(msg)
