"""Templates for workflow steps and workflows, with materialization logic."""

from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Dict, List, Tuple, Any, Callable, Optional, Union

from prefect_cwl.constants import JOBROOT, INROOT
from prefect_cwl.exceptions import ValidationError
from prefect_cwl.backends.artifacts import collect_out_artifacts_local
from prefect_cwl.models import WorkflowStep, CommandLineToolNode, ResourceRequirement


ArtifactPath = Union[Path, List[Path]]


@dataclass(frozen=True)
class StepResources:
    """Normalized runtime resources extracted from CWL ResourceRequirement."""

    cpu_request: Optional[float] = None
    cpu_limit: Optional[float] = None
    memory_request_mb: Optional[int] = None
    memory_limit_mb: Optional[int] = None


def normalize_step_resources(
    resource_req: Optional[ResourceRequirement],
) -> StepResources:
    """Convert CWL ResourceRequirement into executor-oriented resources."""
    if resource_req is None:
        return StepResources()
    return StepResources(
        cpu_request=resource_req.coresMin,
        cpu_limit=resource_req.coresMax,
        memory_request_mb=resource_req.ramMin,
        memory_limit_mb=resource_req.ramMax,
    )


def step_host_dirs(
    workspace: Path, step_name: str, job_suffix: Optional[str] = None
) -> Tuple[Path, Path]:
    """Compute host directories for step output and job directory.

    Args:
        workspace (Path): Workspace root directory
        step_name (str): Step name

    Returns:
         Tuple[Path, Path]: Host directories for step output and job directory
    """
    run_name = f"{step_name}-{job_suffix}" if job_suffix else step_name
    host_outdir = workspace / "steps" / run_name / "out"
    host_jobdir = workspace / "steps" / run_name / "job"
    return host_outdir, host_jobdir


@dataclass(frozen=True)
class ResolvedInputs:
    """Resolved inputs and mounts for a step.

    Attributes:
        volumes: Volume mounts. For Kubernetes, each dict contains mountPath and subPath.
    """

    values: Dict[str, Dict[str, Any]]
    volumes: Dict[str, str]


@dataclass(frozen=True)
class ListingMaterialization:
    """Materialized listing entry to be written on the host filesystem.

    Attributes:
        host_path: Absolute host path where the content will be written.
        content: Text content to write at host_path.
    """

    host_path: Path
    content: str


@dataclass
class StepPlan:
    """Fully materialized step ready for execution.

    Attributes:
        step_name: Unique name used for the container/job.
        tool_id: ID of the CWL tool.
        image: Container image reference.
        argv: Command and arguments to execute.
        outdir_container: Container path for the step output directory.
        volumes: Volume mounts. For Kubernetes, each dict contains mountPath and subPath.
        listings: Files/directories to materialize before execution.
        out_artifacts: Paths to output artifacts produced by the step.
        envs: Environment variables for the container/job.
    """

    step_name: str
    tool_id: str
    image: str
    argv: List[str]
    outdir_container: PurePosixPath
    volumes: Dict[str, str]  # host -> "container:ro|rw"
    listings: List[ListingMaterialization]
    out_artifacts: Dict[str, ArtifactPath]  # tool output name -> host path(s)
    envs: Dict[str, str]
    resources: StepResources = field(default_factory=StepResources)
    host_outdir: Optional[Path] = None
    values: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    container_user: Optional[str] = None


@dataclass
class StepTemplate:
    """Template for a step - no actual input values yet.

    Holds static information parsed from CWL. At runtime, call materialize_step
    to resolve inputs, mounts, argv, and listings.

    Attributes:
        step_name: Unique name used for the container/job.
        tool_id: ID of the CWL tool.
        tool: Parsed CommandLineToolNode.
        wf_step: Parsed WorkflowStep.
        image: Container image reference.
        outdir_container: Container path for the step output directory.
        envs: Environment variables for the container/job.
    """

    step_name: str
    tool_id: str
    tool: CommandLineToolNode
    wf_step: WorkflowStep
    image: str
    outdir_container: PurePosixPath
    envs: Dict[str, str]
    resources: StepResources = field(default_factory=StepResources)
    container_user: Optional[str] = None

    def _compute_base_step_volumes(
        self,
        *,
        host_outdir: Path,
        outdir_container: PurePosixPath,
        host_jobdir: Path,
    ) -> Dict[str, str]:
        return {
            str(host_outdir): f"{outdir_container}:rw",
            str(host_jobdir): f"{JOBROOT}:rw",
        }

    def _resolve_step_inputs_and_mounts(
        self,
        *,
        wf_step: WorkflowStep,
        clt: CommandLineToolNode,
        workflow_inputs: Dict[str, Any],
        produced: Dict[Tuple[str, str], ArtifactPath],
        base_volumes: Dict[str, str],
        input_overrides: Optional[Dict[str, Any]] = None,
    ) -> ResolvedInputs:
        values: Dict[str, Dict[str, Any]] = {
            "workflow": dict(workflow_inputs),
            "inputs": {},
        }
        volumes = dict(base_volumes)

        def as_cwl_directory(container_path: PurePosixPath) -> Dict[str, Any]:
            return {"class": "Directory", "path": str(container_path)}

        def as_cwl_file(container_path: PurePosixPath) -> Dict[str, Any]:
            return {
                "class": "File",
                "path": str(container_path),
                "basename": container_path.name,
            }

        overrides = input_overrides or {}

        for inport, src in wf_step.in_.items():
            if inport in overrides:
                override = overrides[inport]
                tool_input = clt.inputs.get(inport)
                if tool_input is None:
                    raise ValidationError(f"Tool {clt.id!r} has no input {inport!r}")

                if tool_input.type == "Directory" and isinstance(override, Path):
                    mount_target = INROOT / inport
                    volumes[str(override)] = f"{mount_target}:ro"
                    values["inputs"][inport] = as_cwl_directory(mount_target)
                elif tool_input.type == "File" and isinstance(override, Path):
                    mount_target_dir = INROOT / inport
                    volumes[str(override.parent)] = f"{mount_target_dir}:ro"
                    file_in_container = mount_target_dir / override.name
                    values["inputs"][inport] = as_cwl_file(file_in_container)
                elif tool_input.type == "Directory[]" and isinstance(override, list):
                    dirs = []
                    for idx, item in enumerate(override):
                        mount_target = INROOT / inport / str(idx)
                        volumes[str(item)] = f"{mount_target}:ro"
                        dirs.append(as_cwl_directory(mount_target))
                    values["inputs"][inport] = dirs
                elif tool_input.type == "File[]" and isinstance(override, list):
                    files = []
                    for idx, item in enumerate(override):
                        mount_target_dir = INROOT / inport / str(idx)
                        volumes[str(item.parent)] = f"{mount_target_dir}:ro"
                        file_in_container = mount_target_dir / item.name
                        files.append(as_cwl_file(file_in_container))
                    values["inputs"][inport] = files
                else:
                    values["inputs"][inport] = override
                continue

            if isinstance(src, str) and "/" in src:
                upstream_step, upstream_outport = src.split("/", 1)
                host_art = produced.get((upstream_step, upstream_outport))
                if host_art is None:
                    raise ValidationError(
                        f"Step {self.step_name!r} depends on {src!r} but it wasn't produced"
                    )

                tool_input = clt.inputs.get(inport)
                if tool_input is None:
                    raise ValidationError(f"Tool {clt.id!r} has no input {inport!r}")

                if tool_input.type == "Directory":
                    if isinstance(host_art, list):
                        raise ValidationError(
                            f"Input {inport!r} expects Directory but got array from {src!r}"
                        )
                    mount_target = INROOT / inport
                    volumes[str(host_art)] = f"{mount_target}:ro"
                    values["inputs"][inport] = as_cwl_directory(mount_target)

                elif tool_input.type == "File":
                    if isinstance(host_art, list):
                        raise ValidationError(
                            f"Input {inport!r} expects File but got array from {src!r}"
                        )
                    mount_target_dir = INROOT / inport
                    volumes[str(host_art.parent)] = f"{mount_target_dir}:ro"
                    file_in_container = mount_target_dir / host_art.name
                    values["inputs"][inport] = as_cwl_file(file_in_container)

                elif tool_input.type == "Directory[]":
                    if not isinstance(host_art, list):
                        raise ValidationError(
                            f"Input {inport!r} expects Directory[] but got scalar from {src!r}"
                        )
                    dirs = []
                    for idx, item in enumerate(host_art):
                        mount_target = INROOT / inport / str(idx)
                        volumes[str(item)] = f"{mount_target}:ro"
                        dirs.append(as_cwl_directory(mount_target))
                    values["inputs"][inport] = dirs

                elif tool_input.type == "File[]":
                    if not isinstance(host_art, list):
                        raise ValidationError(
                            f"Input {inport!r} expects File[] but got scalar from {src!r}"
                        )
                    files = []
                    for idx, item in enumerate(host_art):
                        mount_target_dir = INROOT / inport / str(idx)
                        volumes[str(item.parent)] = f"{mount_target_dir}:ro"
                        file_in_container = mount_target_dir / item.name
                        files.append(as_cwl_file(file_in_container))
                    values["inputs"][inport] = files

                else:
                    raise ValidationError(
                        f"Input {inport!r} has unsupported type {tool_input.type!r} for upstream wiring"
                    )

            else:
                if src not in workflow_inputs:
                    raise ValidationError(
                        f"Step {self.step_name!r} input {inport!r} references unknown workflow input {src!r}"
                    )
                values["inputs"][inport] = workflow_inputs[src]

        return ResolvedInputs(values=values, volumes=volumes)

    def materialize_step(
        self,
        *,
        workflow_inputs: Dict[str, Any],
        produced: Dict[Tuple[str, str], ArtifactPath],
        workspace: Path,
        render_io,
        job_suffix: Optional[str] = None,
        input_overrides: Optional[Dict[str, Any]] = None,
    ) -> StepPlan:
        """
        Convert a StepTemplate into a StepPlan by resolving actual runtime values.

        This happens at Prefect task execution time, not at plan time.

        Args:
            workflow_inputs: Actual runtime input values
            produced: Map of (step_name, output_port) -> Path for upstream outputs
            workspace: Workspace root directory
            render_io: Function to render command and listings

        Returns:
            Fully materialized StepPlan ready for execution
        """
        host_outdir, host_jobdir = step_host_dirs(
            workspace, self.step_name, job_suffix=job_suffix
        )

        base_volumes = self._compute_base_step_volumes(
            host_outdir=host_outdir,
            outdir_container=self.outdir_container,
            host_jobdir=host_jobdir,
        )

        resolved = self._resolve_step_inputs_and_mounts(
            wf_step=self.wf_step,
            clt=self.tool,
            workflow_inputs=workflow_inputs,
            produced=produced,
            base_volumes=base_volumes,
            input_overrides=input_overrides,
        )

        argv, rendered_listing = render_io(self.tool, resolved.values)

        mats = validate_and_materialize_listings(
            rendered_listing=rendered_listing,
            host_jobdir=host_jobdir,
        )

        out_artifacts = compute_out_artifacts(clt=self.tool, host_outdir=host_outdir)

        return StepPlan(
            step_name=self.step_name,
            tool_id=self.tool_id,
            image=self.image,
            argv=argv,
            outdir_container=self.outdir_container,
            volumes=resolved.volumes,
            listings=mats,
            out_artifacts=out_artifacts,
            envs=self.envs,
            resources=self.resources,
            host_outdir=host_outdir,
            values=resolved.values,
            container_user=self.container_user,
        )


@dataclass
class WorkflowTemplate:
    """Workflow structure without materialized values.

    Attributes:
        workflow_id: ID of the workflow.
        workflow_inputs: Input schemas (types), not runtime values.
        workspace: Workspace root on the host filesystem.
        step_templates: Templates for each workflow step.
        workflow_outputs: Output schemas describing sources.
        waves: Topological waves of step names.
    """

    workflow_id: str
    workflow_inputs: Dict[str, Any]  # Input schemas (types), not values
    workspace: Path
    step_templates: Dict[str, StepTemplate]
    workflow_outputs: Dict[str, Any]  # Output schemas
    waves: List[List[str]]

    def iter_steps(self):
        """Iterate over StepTemplate waves according to topological order."""
        for wave in self.waves:
            yield [self.step_templates.get(step) for step in wave]


def validate_and_materialize_listings(
    *,
    rendered_listing: List[Dict[str, str]],
    host_jobdir: Path,
) -> List[ListingMaterialization]:
    """Validate and materialize InitialWorkDirRequirement listings.

    Args:
        rendered_listing (List[Dict[str, str]]): List of dicts with "entryname" and "entry" keys.
        host_jobdir (Path): Host directory where the listings will be written.
    """
    mats: List[ListingMaterialization] = []

    for item in rendered_listing:
        entryname = item["entryname"]
        entry = item["entry"]

        container_path = PurePosixPath(entryname)

        # Validate under JOBROOT
        try:
            rel = container_path.relative_to(JOBROOT)
        except ValueError:
            raise ValidationError(
                f"Listing entryname must be under {JOBROOT}: got {entryname!r}"
            )

        host_path = host_jobdir / rel
        mats.append(ListingMaterialization(host_path=host_path, content=entry))

    return mats


def compute_out_artifacts(
    *, clt: CommandLineToolNode, host_outdir: Path
) -> Dict[str, ArtifactPath]:
    """Compute output artifacts paths based on output bindings.

    Args:
        clt (CommandLineToolNode): The CommandLineToolNode to compute artifacts for.
        host_outdir (Path): The host directory where the output artifacts will be written.
    """
    out_artifacts: Dict[str, ArtifactPath] = {}
    for outport, outspec in clt.outputs.items():
        glob = outspec.outputBinding.glob
        out_artifacts[outport] = host_outdir / glob
    return out_artifacts


def collect_out_artifacts(
    *,
    clt: CommandLineToolNode,
    host_outdir: Path,
    values: Dict[str, Dict[str, Any]],
) -> Dict[str, ArtifactPath]:
    """Collect step outputs from host filesystem using rendered output globs."""
    return collect_out_artifacts_local(
        clt=clt,
        host_outdir=host_outdir,
        values=values,
    )


RenderIOFn = Callable[
    [CommandLineToolNode, Dict[str, Dict[str, Any]]],
    Tuple[List[str], List[Dict[str, str]]],
]
