"""Decorator for turning Python functions into workflow steps.

Usage::

    from areteos import python_step, InputMapping
    from pydantic import BaseModel

    # Simple — plain dict in/out
    @python_step
    def extract(input: dict) -> dict:
        return {"text": input["raw_text"].upper()}

    step = extract.as_step("extract", "Extract Text")

    # Typed — Pydantic models for validation
    class SummarizeInput(BaseModel):
        text: str
        max_length: int = 100

    class SummarizeOutput(BaseModel):
        summary: str
        word_count: int

    @python_step
    def summarize(input: SummarizeInput) -> SummarizeOutput:
        words = input.text.split()[:input.max_length]
        return SummarizeOutput(summary=" ".join(words), word_count=len(words))

    step = summarize.as_step("summarize", "Summarize")

    # With dependencies (PEP 723 inline metadata for uv run)
    @python_step(requires=["pandas>=2.0", "numpy"])
    def analyze(input: dict) -> dict:
        import pandas as pd
        df = pd.DataFrame(input["data"])
        return {"mean": float(df["value"].mean())}

    step = analyze.as_step("analyze", "Analyze Data")
"""

from __future__ import annotations

import inspect
import textwrap
from collections.abc import Callable
from typing import Any, TypeVar, get_type_hints, overload

from pydantic import BaseModel

from .models import InputMapping, Step, StepType

F = TypeVar("F", bound=Callable[..., Any])


def _is_pydantic_model(annotation: Any) -> bool:
    """Check if an annotation is a Pydantic BaseModel subclass."""
    try:
        return isinstance(annotation, type) and issubclass(annotation, BaseModel)
    except TypeError:
        return False


def _generate_pep723_header(
    requires: list[str],
    python_version: str = "3.11",
) -> str:
    """Generate PEP 723 inline script metadata block."""
    deps = "\n".join(f'#     "{dep}",' for dep in requires)
    return (
        "# /// script\n"
        f'# requires-python = ">={python_version}"\n'
        "# dependencies = [\n"
        f"{deps}\n"
        "# ]\n"
        "# ///\n"
    )


def _generate_script(
    fn: Callable[..., Any],
    requires: list[str] | None,
    python_version: str,
) -> str:
    """Generate the complete inline script from a decorated function."""
    source = inspect.getsource(fn)
    # Remove the decorator line(s) from the source
    source = _strip_decorator(source)
    # Dedent the function body
    source = textwrap.dedent(source)

    # Inspect type hints for input/output models
    hints = get_type_hints(fn)
    input_type = hints.get("input") or hints.get(list(hints.keys())[0] if hints else "input")
    return_type = hints.get("return")

    has_pydantic_input = _is_pydantic_model(input_type)
    has_pydantic_output = _is_pydantic_model(return_type)

    parts: list[str] = []

    # PEP 723 header
    pep723_requires = list(requires or [])
    if (has_pydantic_input or has_pydantic_output) and not any(
        r.startswith("pydantic") for r in pep723_requires
    ):
        pep723_requires.append("pydantic>=2.7")

    if pep723_requires:
        parts.append(_generate_pep723_header(pep723_requires, python_version))

    # Imports
    imports = ["import json", "import sys"]
    if has_pydantic_input or has_pydantic_output:
        imports.append("from pydantic import BaseModel")
    parts.append("\n".join(imports))

    # Pydantic model definitions (inline in the script)
    if has_pydantic_input and input_type is not None:
        parts.append(_generate_model_source(input_type))
    if has_pydantic_output and return_type is not None and return_type is not input_type:
        parts.append(_generate_model_source(return_type))

    # The user's function
    parts.append(source)

    # Main block
    fn_name = fn.__name__
    main_lines = [
        "",
        "def _main():",
        "    data = {}",
        "    if len(sys.argv) > 1:",
        "        with open(sys.argv[1]) as _f:",
        "            data = json.load(_f)",
    ]

    if has_pydantic_input and input_type is not None:
        main_lines.append(f"    _input = {input_type.__name__}.model_validate(data)")
    else:
        main_lines.append("    _input = data")

    main_lines.append(f"    _output = {fn_name}(_input)")

    if has_pydantic_output:
        main_lines.append(
            "    _out = _output.model_dump() if hasattr(_output, 'model_dump') else _output"
        )
        main_lines.append("    print(json.dumps(_out))")
    else:
        main_lines.append("    print(json.dumps(_output))")

    main_lines.extend(
        [
            "",
            "_main()",
        ]
    )

    parts.append("\n".join(main_lines))

    return "\n\n".join(parts) + "\n"


def _strip_decorator(source: str) -> str:
    """Remove @python_step decorator line(s) from function source."""
    lines = source.split("\n")
    result = []
    skip = False
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("@python_step"):
            skip = True
            # Check if decorator continues on next line (multi-line call)
            if "(" in stripped and ")" not in stripped:
                continue
            skip = False
            continue
        if skip:
            if ")" in stripped:
                skip = False
            continue
        result.append(line)
    return "\n".join(result)


def _generate_model_source(model: type[BaseModel]) -> str:
    """Generate Pydantic model class source for embedding in the script."""
    source = inspect.getsource(model)
    return textwrap.dedent(source)


class PythonStepWrapper:
    """Wrapper around a decorated function that can produce Step instances."""

    def __init__(
        self,
        fn: Callable[..., Any],
        requires: list[str] | None = None,
        python_version: str = "3.11",
    ) -> None:
        self._fn = fn
        self._requires = requires
        self._python_version = python_version
        self._script: str | None = None

        # Preserve function metadata
        self.__name__ = fn.__name__
        self.__doc__ = fn.__doc__
        self.__module__ = fn.__module__

        # Extract type hints for validation
        self._hints = get_type_hints(fn)
        self._input_type = None
        self._output_type = self._hints.get("return")

        # Get the first parameter's type hint (the input)
        params = list(inspect.signature(fn).parameters.values())
        if params:
            self._input_type = self._hints.get(params[0].name)

    @property
    def script(self) -> str:
        """The generated inline script content."""
        if self._script is None:
            self._script = _generate_script(self._fn, self._requires, self._python_version)
        return self._script

    @property
    def input_type(self) -> type[BaseModel] | None:
        """The Pydantic input model, if typed."""
        if _is_pydantic_model(self._input_type):
            return self._input_type
        return None

    @property
    def output_type(self) -> type[BaseModel] | None:
        """The Pydantic output model, if typed."""
        if _is_pydantic_model(self._output_type):
            return self._output_type
        return None

    @property
    def input_fields(self) -> list[str]:
        """Field names expected in the input, derived from type hints."""
        if self.input_type is not None:
            return list(self.input_type.model_fields.keys())
        return []

    @property
    def output_fields(self) -> list[str]:
        """Field names produced in the output, derived from type hints."""
        if self.output_type is not None:
            return list(self.output_type.model_fields.keys())
        return []

    def as_step(
        self,
        id: str,
        name: str,
        *,
        input_mappings: list[InputMapping] | None = None,
        timeout_ms: int = 30_000,
        max_retries: int = 1,
    ) -> Step:
        """Create a Step from this decorated function."""
        return Step(
            id=id,
            type=StepType.TOOL,
            name=name,
            config={"tool": "python", "script_content": self.script},
            timeout_ms=timeout_ms,
            max_retries=max_retries,
            input_mappings=input_mappings or [],
        )

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Call the original function directly (for local testing)."""
        return self._fn(*args, **kwargs)


@overload
def python_step(fn: F) -> PythonStepWrapper: ...


@overload
def python_step(
    *,
    requires: list[str] | None = None,
    python_version: str = "3.11",
) -> Callable[[F], PythonStepWrapper]: ...


def python_step(
    fn: F | None = None,
    *,
    requires: list[str] | None = None,
    python_version: str = "3.11",
) -> PythonStepWrapper | Callable[[F], PythonStepWrapper]:
    """Decorator to turn a Python function into a workflow step.

    Can be used with or without arguments::

        @python_step
        def my_step(input: dict) -> dict: ...

        @python_step(requires=["pandas>=2.0"])
        def my_step(input: dict) -> dict: ...

    The decorated function:
    - Receives a dict (or Pydantic model) as its single argument
    - Returns a dict (or Pydantic model) that becomes the step output
    - Can be called directly for local testing
    - Produces a Step via `my_step.as_step("id", "Name")`
    """
    if fn is not None:
        # @python_step without arguments
        return PythonStepWrapper(fn)

    # @python_step(...) with arguments
    def decorator(f: F) -> PythonStepWrapper:
        return PythonStepWrapper(f, requires=requires, python_version=python_version)

    return decorator
