from pathlib import Path
from typing import TYPE_CHECKING

from mthds.models.pipeline_inputs import PipelineInputs

from pipelex import log
from pipelex.config import get_config
from pipelex.core.interpreter.interpreter import PipelexInterpreter
from pipelex.core.memory.working_memory import WorkingMemory
from pipelex.core.memory.working_memory_factory import WorkingMemoryFactory
from pipelex.core.pipes.pipe_factory import PipeFactory
from pipelex.graph.graph_tracer_manager import GraphTracerManager
from pipelex.hub import (
    get_library_manager,
    get_otel_tracer,
    get_pipeline_manager,
    get_report_delegate,
    get_required_pipe,
    get_telemetry_manager,
    resolve_library_dirs,
    set_current_library,
    teardown_current_library,
)
from pipelex.pipe_run.dry_run import convert_to_working_memory_format
from pipelex.pipe_run.pipe_job import PipeJob
from pipelex.pipe_run.pipe_job_factory import PipeJobFactory
from pipelex.pipe_run.pipe_run_mode import PipeRunMode
from pipelex.pipe_run.pipe_run_params import (
    FORCE_DRY_RUN_MODE_ENV_KEY,
    VariableMultiplicity,
)
from pipelex.pipe_run.pipe_run_params_factory import PipeRunParamsFactory
from pipelex.pipeline.exceptions import PipeExecutionError
from pipelex.pipeline.input_normalizer import normalize_data_urls_to_storage
from pipelex.pipeline.job_metadata import JobMetadata, OtelContext
from pipelex.system.configuration.configs import PipelineExecutionConfig
from pipelex.system.environment import get_optional_env
from pipelex.system.telemetry.events import EventName, EventProperty
from pipelex.system.telemetry.otel_constants import OTelConstants
from pipelex.system.telemetry.otel_factory import OtelFactory
from pipelex.tracing.event_log_factory import make_event_log

if TYPE_CHECKING:
    from pipelex.core.bundles.pipelex_bundle_blueprint import PipelexBundleBlueprint
    from pipelex.core.pipes.pipe_abstract import PipeAbstract
    from pipelex.graph.graph_context import GraphContext
    from pipelex.tracing.event_log_protocol import EventLogProtocol


async def pipeline_run_setup(
    execution_config: PipelineExecutionConfig,
    library_id: str | None = None,
    library_dirs: list[str] | None = None,
    pipe_code: str | None = None,
    mthds_contents: list[str] | None = None,
    bundle_uris: list[str] | None = None,
    inputs: PipelineInputs | WorkingMemory | None = None,
    output_name: str | None = None,
    output_multiplicity: VariableMultiplicity | None = None,
    dynamic_output_concept_ref: str | None = None,
    pipe_run_mode: PipeRunMode | None = None,
    search_domain_codes: list[str] | None = None,
    user_id: str | None = None,
    pipeline_run_id: str | None = None,
) -> tuple[PipeJob, str, str]:
    """Set up a pipeline for execution.

    This function handles all the common setup logic for both ``execute_pipeline``
    and ``start_pipeline``, including library setup, pipe loading, working memory
    initialization, and pipe job creation.

    Parameters
    ----------
    execution_config:
        Pipeline execution configuration including graph tracing settings.
        Must be provided by the caller (typically resolved at the entry point).
    library_id:
        Unique identifier for the library instance. If not provided, defaults to the
        auto-generated ``pipeline_run_id``. Use a custom ID when you need to manage
        multiple library instances or maintain library state across executions.
    library_dirs:
        List of directory paths to load pipe definitions from. Combined with directories
        from the ``PIPELEXPATH`` environment variable (PIPELEXPATH directories are searched
        first). When provided alongside ``mthds_contents``, definitions from both sources
        are loaded into the library.
    pipe_code:
        Code identifying the pipe to execute. Required when ``mthds_contents`` is not
        provided. When both ``mthds_contents`` and ``pipe_code`` are provided, the
        specified pipe from the MTHDS contents will be executed (overriding any
        ``main_pipe`` defined in the contents).
    mthds_contents:
        List of MTHDS file contents as strings. The pipe to execute is determined by
        ``pipe_code`` (if provided) or the ``main_pipe`` property in the first content
        that declares one. Can be combined with ``library_dirs`` to load additional
        definitions.
    bundle_uris:
        List of URIs identifying the bundles. Used to detect if bundles were already
        loaded from library directories (e.g., via PIPELEXPATH) to avoid duplicate
        domain registration. If all resolved absolute paths are already in the loaded
        MTHDS paths, the ``mthds_contents`` loading will be skipped.
    inputs:
        Inputs passed to the pipeline. Can be either a ``PipelineInputs`` dictionary
        or a ``WorkingMemory`` instance.
    output_name:
        Name of the output slot to write to.
    output_multiplicity:
        Output multiplicity specification.
    dynamic_output_concept_ref:
        Override the dynamic output concept ref.
    pipe_run_mode:
        Pipe run mode: ``PipeRunMode.LIVE`` or ``PipeRunMode.DRY``. If not specified,
        inferred from the environment variable ``PIPELEX_FORCE_DRY_RUN_MODE``. Defaults
        to ``PipeRunMode.LIVE`` if the environment variable is not set.
    search_domain_codes:
        List of domain codes to search for pipes. The executed pipe's domain is automatically
        added if not already present.
    user_id:
        Unique identifier for the user (optional).
    pipeline_run_id:
        Pre-generated pipeline run ID. If provided, this ID is used instead of
        generating a new one. Use this when the run record has already been created
        externally (e.g., by an API Gateway).

    Returns:
    -------
    tuple[PipeJob, str, str]
        A tuple containing the pipe job ready for execution, the pipeline run ID,
        and the library ID.

    """
    user_id = user_id or OTelConstants.DEFAULT_USER_ID
    if not mthds_contents and not pipe_code:
        msg = "Either pipe_code or mthds_contents must be provided to the pipeline API."
        raise ValueError(msg)

    pipeline = get_pipeline_manager().add_new_pipeline(pipe_code=pipe_code, pipeline_run_id=pipeline_run_id)
    pipeline_run_id = pipeline.pipeline_run_id

    if not library_id:
        library_id = pipeline_run_id

    library_manager = get_library_manager()
    set_current_library(library_id=library_id)
    library_manager.open_library(library_id=library_id)

    pipe: PipeAbstract | None = None
    blueprint: PipelexBundleBlueprint | None = None

    effective_dirs, source_label = resolve_library_dirs(library_dirs)

    if effective_dirs:
        log.verbose(f"Loading libraries from {len(effective_dirs)} directory(ies) ({source_label}):")
        for index_dir, dir_path in enumerate(effective_dirs):
            log.verbose(f"  [{index_dir + 1}] {dir_path}")
        library_manager.load_libraries(
            library_id=library_id,
            library_dirs=effective_dirs,
        )
    else:
        log.verbose(f"No library directories to load ({source_label})")

    # Then handle MTHDS content(s) or pipe_code
    if mthds_contents:
        all_blueprints = [PipelexInterpreter.make_pipelex_bundle_blueprint(mthds_content=content) for content in mthds_contents]

        # Filter out blueprints whose URIs are already loaded
        blueprints_to_load: list[PipelexBundleBlueprint] = list(all_blueprints)
        if bundle_uris:
            current_library = library_manager.get_library(library_id=library_id)
            blueprints_to_load = []
            for blueprint, uri in zip(all_blueprints, bundle_uris, strict=True):
                try:
                    resolved_uri = Path(uri).resolve()
                except (OSError, RuntimeError):
                    resolved_uri = Path(uri)
                if resolved_uri in current_library.loaded_mthds_paths:
                    log.verbose(f"Bundle '{uri}' already loaded from library directories, skipping")
                else:
                    blueprints_to_load.append(blueprint)

        if blueprints_to_load:
            library_manager.load_from_blueprints(library_id=library_id, blueprints=blueprints_to_load)

        # Find the pipe to execute
        if pipe_code:
            pipe = get_required_pipe(pipe_code=pipe_code)
        else:
            # Find main_pipe from the first blueprint that declares one
            # Qualify main_pipe with domain to avoid ambiguity when multiple domains define pipes with the same code.
            # Note: main_pipe is validated as snake_case (no dots allowed), so it is always a bare code — never
            # already domain-qualified. See PipelexBundleBlueprint.validate_main_pipe_syntax.
            qualified_main_pipe: str | None = None
            for blueprint in all_blueprints:
                if blueprint.main_pipe:
                    qualified_main_pipe = PipeFactory.make_pipe_ref_with_domain(domain_code=blueprint.domain, pipe_code=blueprint.main_pipe)
                    break
            if not qualified_main_pipe:
                msg = "No pipe_code provided and no main_pipe found in any of the MTHDS contents."
                raise PipeExecutionError(message=msg)
            pipe = get_required_pipe(pipe_code=qualified_main_pipe)

    elif pipe_code:
        pipe = get_required_pipe(pipe_code=pipe_code)
    else:
        msg = "Either provide pipe_code or mthds_contents to the pipeline API."
        raise PipeExecutionError(message=msg)

    pipe_code = pipe.code

    search_domain_codes = search_domain_codes or []
    if pipe.domain_code not in search_domain_codes:
        search_domain_codes.insert(0, pipe.domain_code)

    # Initialize graph tracing if requested (after pipe is loaded so we have domain info)
    graph_context: GraphContext | None = None
    event_log: EventLogProtocol | None = None
    if execution_config.is_generate_graph:
        # Create event log when tracing is enabled
        config = get_config()
        tracing_config = config.pipelex.tracing_config
        if tracing_config.is_enabled:
            event_log = make_event_log(tracing_config)

        graph_tracer_manager = GraphTracerManager.get_or_create_instance()
        graph_context = graph_tracer_manager.open_tracer(
            graph_id=pipeline_run_id,
            data_inclusion=execution_config.graph_config.data_inclusion,
            pipeline_ref_domain=pipe.domain_code,
            pipeline_ref_main_pipe=pipe_code,
            event_log=event_log,
            workflow_id="direct",
            pipeline_run_id=pipeline_run_id,
        )

    try:
        working_memory: WorkingMemory | None = None

        # First, process user-provided inputs
        if inputs:
            if isinstance(inputs, WorkingMemory):
                working_memory = inputs
            else:
                working_memory = WorkingMemoryFactory.make_from_pipeline_inputs(
                    pipeline_inputs=inputs,
                    search_domain_codes=search_domain_codes,
                )

        # If mock inputs is enabled, generate mock data for missing required inputs
        if execution_config.is_mock_inputs:
            needed_inputs_spec = pipe.needed_inputs()
            needed_inputs_for_factory = convert_to_working_memory_format(needed_inputs_spec)

            # Filter out inputs that were already provided by the user
            if working_memory:
                provided_names = set(working_memory.root.keys())
                missing_inputs = [spec for spec in needed_inputs_for_factory if spec.variable_name not in provided_names]
            else:
                missing_inputs = needed_inputs_for_factory
                working_memory = WorkingMemoryFactory.make_empty()

            # Generate mock data only for missing inputs
            if missing_inputs:
                mock_memory = WorkingMemoryFactory.make_mock_inputs(needed_inputs=missing_inputs)
                for name, stuff in mock_memory.root.items():
                    working_memory.add_new_stuff(name=name, stuff=stuff)

        # Normalize data URLs to pipelex-storage:// URIs if configured
        if working_memory and execution_config.is_normalize_data_urls_to_storage and not execution_config.is_mock_inputs:
            working_memory = await normalize_data_urls_to_storage(working_memory)

        # TODO: rethink this, it's not forcing
        if pipe_run_mode is None:
            if run_mode_from_env := get_optional_env(key=FORCE_DRY_RUN_MODE_ENV_KEY):
                pipe_run_mode = PipeRunMode(run_mode_from_env)
            else:
                pipe_run_mode = PipeRunMode.LIVE

        get_report_delegate().open_registry(pipeline_run_id=pipeline_run_id)

        # Set event log on the report delegate for distributed usage event emission
        if event_log is not None:
            get_report_delegate().set_event_log(
                context_key=pipeline_run_id,
                event_log=event_log,
                workflow_id="direct",
                pipeline_run_id=pipeline_run_id,
            )

        # Initialize OtelContext if telemetry is enabled (not dry mode and tracer available)
        # The trace_id is computed once here; span_id uses OTEL_VIRTUAL_ROOT_PARENT_SPAN_ID for root
        otel_context: OtelContext | None = None
        if pipe_run_mode.is_live and get_otel_tracer() is not None:
            trace_id = OtelFactory.make_trace_id(pipeline_run_id=pipeline_run_id)
            trace_name, trace_name_redacted = OtelFactory.make_trace_names(pipeline_run_id=pipeline_run_id, pipe_code=pipe_code)
            otel_context = OtelContext(
                trace_id=trace_id,
                trace_name=trace_name,
                trace_name_redacted=trace_name_redacted,
                span_id=OTelConstants.OTEL_VIRTUAL_ROOT_PARENT_SPAN_ID,
            )
            # Emit trace start event immediately to establish trace name in PostHog
            # This must happen before any pipe spans are created/exported
            get_telemetry_manager().handle_trace_start(trace_name=trace_name, trace_name_redacted=trace_name_redacted, trace_id=trace_id)

        job_metadata = JobMetadata(
            user_id=user_id,
            pipeline_run_id=pipeline.pipeline_run_id,
            otel_context=otel_context,
            graph_context=graph_context,
        )

        pipe_run_params = PipeRunParamsFactory.make_run_params(
            output_multiplicity=output_multiplicity,
            dynamic_output_concept_ref=dynamic_output_concept_ref,
            pipe_run_mode=pipe_run_mode,
        )

        # Build the library crate from all accumulated blueprints for Temporal dispatch
        library_crate = library_manager.get_crate(library_id=library_id)

        pipe_job = PipeJobFactory.make_pipe_job(
            pipe=pipe,
            pipe_run_params=pipe_run_params,
            job_metadata=job_metadata,
            working_memory=working_memory,
            output_name=output_name,
            library_crate=library_crate,
        )

        properties = {
            EventProperty.PIPELINE_RUN_ID: job_metadata.pipeline_run_id,
            EventProperty.PIPE_TYPE: pipe.pipe_type,
        }
        get_telemetry_manager().track_event(event_name=EventName.PIPELINE_EXECUTE, properties=properties)

        return pipe_job, pipeline_run_id, library_id
    except Exception:
        # Error-path-only cleanup: on any failure during setup, tear down the graph tracer, event-log state
        # and library that were partially created, then re-raise. Re-raises — never swallows. Cannot be a
        # `finally`: on success these resources are returned alive for the pipeline run to use.
        # Cleanup graph tracer if it was opened
        if graph_context is not None:
            tracer_manager = GraphTracerManager.get_instance()
            if tracer_manager is not None:
                tracer_manager.close_tracer(pipeline_run_id)
        # Cleanup event log state from the report delegate
        if event_log is not None:
            get_report_delegate().clear_event_log(context_key=pipeline_run_id)
        # Cleanup library
        library_manager.teardown(library_id=library_id)
        teardown_current_library()
        raise
