from pathlib import Path
from typing import Any, Dict, List, Optional

from fastmcp import Client, FastMCP
from fastmcp.client.transports import FastMCPTransport
from satif_core.code_executors.base import CodeExecutor
from satif_core.transformation_builders.base import AsyncTransformationBuilder
from satif_core.types import (
    FilePath,
    SDIFPath,
    TransformationResult,
)
from satif_sdk.code_executors.local_executor import LocalCodeExecutor
from satif_sdk.transformers.code import CodeTransformer
from sdif_mcp.server import mcp

from satif_ai.transformation_builders.syncpulse import SyncpulseTransformationBuilder
from satif_ai.utils.openai_mcp import OpenAICompatibleMCP


async def atransform(
    sdif: SDIFPath,
    output_target_files: Dict[FilePath, str] | List[FilePath] | FilePath,
    instructions: Optional[str] = None,
    output_path: FilePath = Path("."),
    *,
    transformation_code: Optional[str] = None,
    transformation_builder: Optional[AsyncTransformationBuilder] = None,
    code_executor: Optional[CodeExecutor] = None,
    mcp_server: Optional[FastMCP] = None,
    mcp_client: Optional[Client] = None,
    llm_model: str = "o4-mini",
    schema_only: bool = False,
    representer_kwargs: Optional[Dict[str, Any]] = None,
) -> TransformationResult:
    """
    Asynchronously transforms an SDIF (Standard Data Interchange Format) input using
    an AI-generated or provided transformation code.

    This function orchestrates the process of:
    1. Optionally generating transformation code using an AI model via a `TransformationBuilder`
       if `transformation_code` is not provided.
       explicitly passed.
    2. Executing the transformation code using a `CodeTransformer` and a `CodeExecutor`.
    3. Exporting the results to the specified output.

    Args:
        sdif: Path to the input SDIF file or an `SDIFDatabase` object.
        output_target_files: A dictionary mapping original example file paths (or string identifiers)
                             to their desired agent-facing filenames, or a list of output example
                             file paths, or a single output file path. These are used by the AI to understand the target
                             format and structure, and also by the `CodeTransformer` to determine
                             output filenames if the transformation result keys match.
        instructions: Optional. Natural language instructions for the AI to generate
                      the transformation code. Used if `transformation_code` is None.
        output_path: Path to the directory where transformation outputs will be saved.
        transformation_code: Optional. Pre-existing Python code for the transformation.
                             If None, code will be generated by the `transformation_builder`.
        transformation_builder: Optional. An `AsyncTransformationBuilder` instance responsible for generating
                      the transformation code if `transformation_code` is not provided.
                      If None, a `SyncpulseTransformationBuilder` is instantiated.
        code_executor: Optional. A `CodeExecutor` instance for running the transformation
                       code. If None, a `LocalCodeExecutor` is used.
        mcp_server: Optional. A `FastMCP` server instance for the AI code builder.
                    Defaults to the global `mcp` instance if `transformation_builder` is None and
                    a new `SyncpulseTransformationBuilder` is being created.
        mcp_client: Optional. A user-provided `Client` instance. If provided when
                    `transformation_builder` is None, it will be used by the internally
                    created `SyncpulseTransformationBuilder`. The caller is responsible for
                    managing the lifecycle of a provided client.
        llm_model: The language model to use for code generation (e.g., "o4-mini").
                   Used if `transformation_builder` is None.
        schema_only: If True, the transformation aims to match only the schema (headers)
                     of the `output_target_files`, and input samples may be omitted or marked
                     as empty for the AI. This is useful for structural transformations
                     without processing actual data rows.
        representer_kwargs: Optional dictionary of keyword arguments to pass to the
                            representer when analyzing `output_target_files`.

    Returns:
        A `TransformationResult` object containing the path to the output
        and the transformation code used.
    """
    current_transformation_code: Optional[str] = transformation_code
    active_builder: Optional[AsyncTransformationBuilder] = transformation_builder

    _openai_mcp_instance: Optional[OpenAICompatibleMCP] = None
    openai_mcp_managed_locally = False

    # If code isn't provided, we need a builder. If a builder isn't provided, we create one.
    if current_transformation_code is None:
        if active_builder is None:
            _effective_mcp_server = mcp_server if mcp_server is not None else mcp

            _openai_mcp_instance = OpenAICompatibleMCP(mcp=_effective_mcp_server)
            await _openai_mcp_instance.connect()
            openai_mcp_managed_locally = True

            if mcp_client is None:  # No user-provided client, create and manage one
                mcp_transport = FastMCPTransport(mcp=_effective_mcp_server)
                async with Client(mcp_transport) as new_client:
                    active_builder = SyncpulseTransformationBuilder(
                        mcp_server=_openai_mcp_instance,
                        mcp_session=new_client.session,
                        llm_model=llm_model,
                    )
                    current_transformation_code = await active_builder.build(
                        sdif=sdif,
                        output_target_files=output_target_files,
                        instructions=instructions,
                        schema_only=schema_only,
                        representer_kwargs=representer_kwargs,
                    )
            else:
                active_builder = SyncpulseTransformationBuilder(
                    mcp_server=_openai_mcp_instance,
                    mcp_session=mcp_client,  # Use the provided client
                    llm_model=llm_model,
                )
                current_transformation_code = await active_builder.build(
                    sdif=sdif,
                    output_target_files=output_target_files,
                    instructions=instructions,
                    schema_only=schema_only,
                    representer_kwargs=representer_kwargs,
                )

    # Disconnect OpenAICompatibleMCP if it was created and connected locally
    if (
        openai_mcp_managed_locally
        and _openai_mcp_instance
        and _openai_mcp_instance._is_connected
    ):
        await _openai_mcp_instance.cleanup()

    if current_transformation_code is None:
        raise ValueError("Transformation code could not be obtained or generated.")

    _code_executor = code_executor if code_executor is not None else LocalCodeExecutor()

    transformer = CodeTransformer(
        function=current_transformation_code,
        code_executor=_code_executor,
    )

    exported_artifact_path = transformer.export(
        sdif=sdif,
        output_path=output_path,
    )

    return TransformationResult(
        output_path=exported_artifact_path, function_code=current_transformation_code
    )
