from collections.abc import Callable, Generator
from pathlib import Path
from typing import Any, Literal

import mlx.core as mx
from mflux.models.common.config.config import Config
from PIL import Image

from exo.api.types import AdvancedImageParams
from exo.download.download_utils import build_model_path
from exo.shared.types.common import ModelId
from exo.shared.types.worker.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
    ShardMetadata,
)
from exo.worker.engines.image.config import ImageModelConfig
from exo.worker.engines.image.models import (
    create_adapter_for_model,
    get_config_for_model,
)
from exo.worker.engines.image.models.base import ModelAdapter
from exo.worker.engines.image.pipeline import DiffusionRunner
from exo.worker.engines.mlx.utils_mlx import mx_barrier
from exo.worker.runner.bootstrap import logger


class DistributedImageModel:
    model_id: ModelId
    _config: ImageModelConfig
    _adapter: ModelAdapter[Any, Any]
    _runner: DiffusionRunner

    def __init__(
        self,
        model_id: ModelId,
        local_path: Path,
        shard_metadata: PipelineShardMetadata | CfgShardMetadata,
        group: mx.distributed.Group | None,
        quantize: int | None = None,
    ):
        config = get_config_for_model(model_id)
        adapter = create_adapter_for_model(config, model_id, local_path, quantize)

        has_layer_sharding = (
            shard_metadata.start_layer != 0
            or shard_metadata.end_layer != shard_metadata.n_layers
        )

        if group is not None and has_layer_sharding:
            adapter.slice_transformer_blocks(
                start_layer=shard_metadata.start_layer,
                end_layer=shard_metadata.end_layer,
            )

        runner = DiffusionRunner(
            config=config,
            adapter=adapter,
            group=group,
            shard_metadata=shard_metadata,
        )

        if group is not None:
            logger.info("Initialized distributed diffusion runner")

            mx.eval(adapter.model.parameters())  # pyright: ignore[reportAny]

            # TODO(ciaran): Do we need this?
            mx.eval(adapter.model)  # pyright: ignore[reportAny]

            mx_barrier(group)
            logger.info(f"Transformer sharded for rank {group.rank()}")
        else:
            logger.info("Single-node initialization")

        self.model_id = model_id
        self._config = config
        self._adapter = adapter
        self._runner = runner

    @classmethod
    def from_shard_metadata(
        cls, shard: ShardMetadata, group: mx.distributed.Group | None
    ) -> "DistributedImageModel":
        model_id = shard.model_card.model_id
        model_path = build_model_path(model_id)

        if not isinstance(shard, (PipelineShardMetadata, CfgShardMetadata)):
            raise ValueError(
                "Expected PipelineShardMetadata or CfgShardMetadata for image generation"
            )

        return cls(
            model_id=model_id,
            local_path=model_path,
            shard_metadata=shard,
            group=group,
        )

    def get_steps_for_quality(self, quality: Literal["low", "medium", "high"]) -> int:
        """Get the number of inference steps for a quality level."""
        return self._config.get_steps_for_quality(quality)

    def generate(
        self,
        prompt: str,
        height: int,
        width: int,
        quality: Literal["low", "medium", "high"] = "medium",
        seed: int = 2,
        image_path: Path | None = None,
        partial_images: int = 0,
        advanced_params: AdvancedImageParams | None = None,
        cancel_checker: Callable[[], bool] | None = None,
    ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
        if (
            advanced_params is not None
            and advanced_params.num_inference_steps is not None
        ):
            steps = advanced_params.num_inference_steps
        else:
            steps = self._config.get_steps_for_quality(quality)

        guidance_override: float | None = None
        if advanced_params is not None and advanced_params.guidance is not None:
            guidance_override = advanced_params.guidance

        negative_prompt: str | None = None
        if advanced_params is not None and advanced_params.negative_prompt is not None:
            negative_prompt = advanced_params.negative_prompt

        # For edit mode: compute dimensions from input image
        # This also stores image_paths in the adapter for encode_prompt()
        if image_path is not None:
            computed_dims = self._adapter.set_image_dimensions(image_path)
            if computed_dims is not None:
                # Override user-provided dimensions with computed ones
                width, height = computed_dims

        config = Config(
            num_inference_steps=steps,
            height=height,
            width=width,
            image_path=image_path,
            model_config=self._adapter.model.model_config,  # pyright: ignore[reportAny]
            guidance=guidance_override if guidance_override is not None else 4.0,
        )

        if advanced_params is not None and advanced_params.num_sync_steps is not None:
            num_sync_steps = advanced_params.num_sync_steps
        else:
            num_sync_steps = self._config.num_sync_steps

        for result in self._runner.generate_image(
            runtime_config=config,
            prompt=prompt,
            seed=seed,
            partial_images=partial_images,
            guidance_override=guidance_override,
            negative_prompt=negative_prompt,
            num_sync_steps=num_sync_steps,
            cancel_checker=cancel_checker,
        ):
            if isinstance(result, tuple):
                # Partial image: (GeneratedImage, partial_index, total_partials)
                image, partial_idx, total_partials = result
                yield (image, partial_idx, total_partials)
            else:
                logger.info("generated image")
                yield result
