import asyncio
from typing import TYPE_CHECKING, Any, Literal, cast

import shortuuid
from typing_extensions import override

from pipelex.config import get_config
from pipelex.core.memory.working_memory import MAIN_STUFF_NAME, WorkingMemory
from pipelex.core.pipes.inputs.input_requirements import InputRequirements
from pipelex.core.pipes.pipe_output import PipeOutput
from pipelex.core.stuffs.list_content import ListContent
from pipelex.core.stuffs.stuff_factory import StuffFactory
from pipelex.hub import get_pipeline_tracker, get_required_pipe
from pipelex.pipe_controllers.pipe_controller import PipeController
from pipelex.pipe_run.exceptions import PipeRunError
from pipelex.pipe_run.pipe_run_params import BatchParams, PipeRunMode, PipeRunParams
from pipelex.pipeline.job_metadata import JobMetadata

if TYPE_CHECKING:
    from collections.abc import Coroutine

    from pipelex.core.stuffs.stuff import Stuff
    from pipelex.core.stuffs.stuff_content import StuffContent


class PipeBatch(PipeController):
    type: Literal["PipeBatch"] = "PipeBatch"

    branch_pipe_code: str
    batch_params: BatchParams

    @override
    def required_variables(self) -> set[str]:
        required_variables: set[str] = set()
        # 1. Check that the inputs of the pipe branch_pipe_code are in the inputs of the pipe
        pipe = get_required_pipe(pipe_code=self.branch_pipe_code)
        for variable_name, _ in pipe.inputs.items:
            required_variables.add(variable_name)
        # 2. Check that the input_list_stuff_name is in the inputs of the pipe
        required_variables.remove(self.batch_params.input_item_stuff_name)
        required_variables.add(self.batch_params.input_list_stuff_name)
        return required_variables

    @override
    def needed_inputs(self, visited_pipes: set[str] | None = None) -> InputRequirements:
        return self.inputs

    @override
    def pipe_dependencies(self) -> set[str]:
        return {self.branch_pipe_code}

    @override
    def validate_inputs_static(self):
        pass

    @override
    def validate_inputs_with_library(self):
        pass

    @override
    def validate_output_static(self):
        pass

    @override
    def validate_output_with_library(self):
        pass

    @override
    async def validate_before_run(
        self, job_metadata: JobMetadata, working_memory: WorkingMemory, pipe_run_params: PipeRunParams, output_name: str | None = None
    ) -> None:
        batch_params = pipe_run_params.batch_params or self.batch_params or BatchParams.make_default()
        input_list_stuff_name = batch_params.input_list_stuff_name
        if not self.inputs.is_variable_existing(variable_name=input_list_stuff_name):
            msg = f"Batch input list named '{input_list_stuff_name}' is not in PipeBatch '{self.code}' input requirements: {self.inputs}"
            raise PipeRunError(message=msg, run_mode=pipe_run_params.run_mode)

        if not working_memory.is_stuff_exists(input_list_stuff_name):
            msg = f"Input list stuff '{input_list_stuff_name}' required by this PipeBatch '{self.code}' not found in working memory"
            raise PipeRunError(message=msg, run_mode=pipe_run_params.run_mode)

        input_stuff = working_memory.get_stuff(input_list_stuff_name)
        if not isinstance(input_stuff.content, ListContent):
            msg = (
                f"Input list stuff '{input_list_stuff_name}' of PipeBatch '{self.code}' must be ListContent, "
                f"got {input_stuff.stuff_name or 'unnamed'} = {type(input_stuff.content)}. stuff: {input_stuff}"
            )
            raise PipeRunError(message=msg, run_mode=pipe_run_params.run_mode)

    async def _run_batch_pipe(
        self,
        job_metadata: JobMetadata,
        working_memory: WorkingMemory,
        pipe_run_params: PipeRunParams,
        output_name: str | None = None,
    ) -> PipeOutput:
        """Common logic for running or dry-running a pipe in batch mode."""
        batch_params = pipe_run_params.batch_params or self.batch_params or BatchParams.make_default()
        input_item_stuff_name = batch_params.input_item_stuff_name
        input_list_stuff_name = batch_params.input_list_stuff_name

        if pipe_run_params.final_stuff_code:
            method_name = "dry_run_pipe" if pipe_run_params.run_mode == PipeRunMode.DRY else "_run_controller_pipe"
            pipe_run_params.final_stuff_code = None

        pipe_run_params.push_pipe_layer(pipe_code=self.branch_pipe_code)

        input_stuff = working_memory.get_stuff(input_list_stuff_name)
        input_content = cast("ListContent[StuffContent]", input_stuff.content)

        # TODO: Make commented code work when inputing images named "a.b.c"
        sub_pipe = get_required_pipe(pipe_code=self.branch_pipe_code)
        nb_history_items_limit = get_config().pipelex.tracker_config.applied_nb_items_limit
        batch_output_stuff_code = shortuuid.uuid()
        tasks: list[Coroutine[Any, Any, PipeOutput]] = []
        item_stuffs: list[Stuff] = []
        required_stuff_lists: list[list[Stuff]] = []
        branch_output_item_codes: list[str] = []

        for branch_index, item in enumerate(input_content.items):
            branch_output_item_code = f"{batch_output_stuff_code}-branch-{branch_index}"
            branch_output_item_codes.append(branch_output_item_code)
            if nb_history_items_limit and branch_index >= nb_history_items_limit:
                break
            branch_input_item_code = f"{input_stuff.stuff_code}-branch-{branch_index}"
            item_input_stuff = StuffFactory.make_stuff(
                code=branch_input_item_code,
                concept=self.inputs.get_required_input_requirement(input_list_stuff_name).concept,
                content=item,
                name=input_item_stuff_name,
            )
            item_stuffs.append(item_input_stuff)
            branch_memory = working_memory.make_deep_copy()
            branch_memory.set_new_main_stuff(stuff=item_input_stuff, name=input_item_stuff_name)

            required_variables = sub_pipe.required_variables()
            required_stuffs = branch_memory.get_existing_stuffs(names=required_variables)
            required_stuffs = [required_stuff for required_stuff in required_stuffs if required_stuff.stuff_code != input_stuff.stuff_code]
            required_stuff_lists.append(required_stuffs)
            branch_pipe_run_params = pipe_run_params.deep_copy_with_final_stuff_code(final_stuff_code=branch_output_item_code)

            task: Coroutine[Any, Any, PipeOutput]
            if pipe_run_params.run_mode == PipeRunMode.DRY:
                branch_pipe_run_params.run_mode = PipeRunMode.DRY
                task = sub_pipe.run_pipe(
                    job_metadata=job_metadata,
                    working_memory=branch_memory,
                    output_name=f"Batch result {branch_index + 1} of {output_name}",
                    pipe_run_params=branch_pipe_run_params,
                )
            else:
                task = sub_pipe.run_pipe(
                    job_metadata=job_metadata,
                    working_memory=branch_memory,
                    output_name=f"Batch result {branch_index + 1} of {output_name}",
                    pipe_run_params=branch_pipe_run_params,
                )
            tasks.append(task)

        pipe_outputs = await asyncio.gather(*tasks)

        output_items: list[StuffContent] = []
        output_stuffs: list[Stuff] = []
        output_stuff_code = shortuuid.uuid()[:5]
        for pipe_output in pipe_outputs:
            branch_output_stuff = pipe_output.main_stuff
            output_stuffs.append(branch_output_stuff)
            output_items.append(branch_output_stuff.content)

        list_content: ListContent[StuffContent] = ListContent(items=output_items)
        output_stuff = StuffFactory.make_stuff(
            code=output_stuff_code,
            concept=self.output,
            content=list_content,
            name=output_name,
        )

        method_name = "dry_run_pipe" if pipe_run_params.run_mode == PipeRunMode.DRY else "run_pipe"
        for branch_index, (
            required_stuff_list,
            item_input_stuff,
            item_output_stuff,
        ) in enumerate(zip(required_stuff_lists, item_stuffs, output_stuffs, strict=False)):
            get_pipeline_tracker().add_batch_step(
                from_stuff=input_stuff,
                to_stuff=item_input_stuff,
                to_branch_index=branch_index,
                pipe_layer=pipe_run_params.pipe_layers,
                comment=f"PipeBatch.{method_name}() in zip",
            )
            for required_stuff in required_stuff_list:
                get_pipeline_tracker().add_pipe_step(
                    from_stuff=required_stuff,
                    to_stuff=item_output_stuff,
                    pipe_code=self.branch_pipe_code,
                    pipe_layer=pipe_run_params.pipe_layers,
                    comment=f"PipeBatch.{method_name}() on required_stuff_list",
                    as_item_index=branch_index,
                    is_with_edge=(required_stuff.stuff_name != MAIN_STUFF_NAME),
                )

        for branch_index, branch_output_stuff in enumerate(output_stuffs):
            branch_output_item_code = branch_output_item_codes[branch_index]
            get_pipeline_tracker().add_aggregate_step(
                from_stuff=branch_output_stuff,
                to_stuff=output_stuff,
                pipe_layer=pipe_run_params.pipe_layers,
                comment=f"PipeBatch.{method_name}() on branch_index of batch",
            )

        working_memory.set_new_main_stuff(
            stuff=output_stuff,
            name=output_name,
        )

        return PipeOutput(
            working_memory=working_memory,
            pipeline_run_id=job_metadata.pipeline_run_id,
        )

    @override
    async def _run_controller_pipe(
        self,
        job_metadata: JobMetadata,
        working_memory: WorkingMemory,
        pipe_run_params: PipeRunParams,
        output_name: str | None = None,
    ) -> PipeOutput:
        return await self._run_batch_pipe(
            job_metadata=job_metadata,
            working_memory=working_memory,
            pipe_run_params=pipe_run_params,
            output_name=output_name,
        )

    @override
    async def _dry_run_controller_pipe(
        self,
        job_metadata: JobMetadata,
        working_memory: WorkingMemory,
        pipe_run_params: PipeRunParams,
        output_name: str | None = None,
    ) -> PipeOutput:
        return await self._run_batch_pipe(
            job_metadata=job_metadata,
            working_memory=working_memory,
            pipe_run_params=pipe_run_params,
            output_name=output_name,
        )
