import inspect
import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable

from jinja2 import Environment, FileSystemLoader, Template, select_autoescape
from jsonschema import ValidationError, validate

from artificer_workflows.artifacts import Artifact, create_artifact

if TYPE_CHECKING:
    from fastmcp import FastMCP


DEFAULT_STEP_TEMPLATE = """\
Workflow: {{ workflow_name }}
Execution ID: {{ execution_id }}
Step: {{ step_name }}
Step ID: {{ step_id }}

## Instructions
{{ instructions }}

## Result Schema
{{ result_schema }}

## Next Action
When complete, call `{{ workflow_name }}__complete_step` with:
  - execution_id: {{ execution_id }}
  - step_id: {{ step_id }}
  - status: SUCCESS or ERROR
  - output: <your output matching the schema above>

Use status=SUCCESS when the step completed successfully.
Use status=ERROR if you encountered an error - the step will be retried.
"""


class Workflow:
    """Base class for workflows. Subclass to create a workflow."""

    _tools: list[tuple[str, Callable]]
    _steps: dict[str, type]  # step class name -> step class
    _start_step: str | None
    _active_executions: dict[str, object]  # execution_id -> current step instance
    _output_dirs: dict[str, Path]  # execution_id -> output directory
    _artifacts: dict[str, list[Artifact]]  # execution_id -> list of artifacts
    _jinja_env: Environment | None  # Jinja2 environment for templates
    Step: type  # Step base class for this workflow
    output_dir: Path = Path(".artificer")  # base output directory
    templates_dir: Path | None = None  # Override to set templates directory
    step_template: str | None = None  # Override to load step template from file

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # Initialize class-level state
        cls._tools = []
        cls._steps = {}
        cls._start_step = None
        cls._active_executions = {}
        cls._output_dirs = {}
        cls._artifacts = {}

        # Set up Jinja2 environment if templates_dir is set
        if cls.templates_dir is not None:
            cls._jinja_env = Environment(
                loader=FileSystemLoader(cls.templates_dir),
                autoescape=select_autoescape(),
            )
        else:
            cls._jinja_env = None

        cls.Step = cls._create_step_class()

    @classmethod
    def _create_step_class(cls):
        workflow = cls

        class Step:
            """Base class for workflow steps.

            Define instructions_template and result_schema.
            """

            # Template filename (e.g., "collect_requirements.md")
            instructions_template: str
            # JSON Schema for step output
            result_schema: dict[str, Any] = {"type": "object"}
            max_retries: int = 3  # Can be overridden per step
            id: str
            execution_id: str | None
            result: dict[str, Any] | None  # Output from previous step

            def __init__(
                self,
                execution_id: str | None = None,
                result: dict[str, Any] | None = None,
            ):
                self.id = f"{uuid.uuid4()}:attempt:1"
                self.execution_id = execution_id
                self.result = result

            @staticmethod
            def _parse_step_id(step_id: str) -> tuple[str, int]:
                """Parse step ID into (uuid, attempt)."""
                parts = step_id.rsplit(":attempt:", 1)
                if len(parts) == 2:
                    return parts[0], int(parts[1])
                return step_id, 1

            @staticmethod
            def _make_step_id(step_uuid: str, attempt: int) -> str:
                """Create step ID from uuid and attempt."""
                return f"{step_uuid}:attempt:{attempt}"

            def instructions(self) -> str:
                """Render the instructions template with context."""
                if workflow._jinja_env is None:
                    msg = f"Workflow '{workflow.__name__}' has no templates_dir set"
                    raise RuntimeError(msg)
                template = workflow._jinja_env.get_template(self.instructions_template)
                return template.render(result=self.result)

            def render(self) -> str:
                """Render the full step prompt."""
                # Use file template if specified, otherwise use inline default
                if workflow.step_template is not None:
                    if workflow._jinja_env is None:
                        msg = f"Workflow '{workflow.__name__}' has no templates_dir set"
                        raise RuntimeError(msg)
                    template = workflow._jinja_env.get_template(workflow.step_template)
                else:
                    template = Template(DEFAULT_STEP_TEMPLATE)
                return template.render(
                    workflow_name=workflow.__name__,
                    execution_id=self.execution_id or "N/A",
                    step_name=type(self).__name__,
                    step_id=self.id,
                    instructions=self.instructions(),
                    result_schema=json.dumps(self.result_schema, indent=2),
                )

            def __init_subclass__(step_cls, start: bool = False, **kwargs):
                super().__init_subclass__(**kwargs)
                if not hasattr(step_cls, "instructions_template"):
                    msg = f"{step_cls.__name__} must define 'instructions_template'"
                    raise TypeError(msg)

                # Store class by name
                step_name = step_cls.__name__
                workflow._steps[step_name] = step_cls

                # Track start step
                if start:
                    if workflow._start_step is not None:
                        msg = f"Workflow already has start step: {workflow._start_step}"
                        raise TypeError(msg)
                    workflow._start_step = step_name

                # Collect for tool registration (wrapper creates instance per call)
                def make_tool(cls):
                    def tool_fn():
                        instance = cls()
                        return instance.instructions()

                    tool_fn.__doc__ = f"Get instructions for {cls.__name__}"
                    return tool_fn

                tool_name = f"{workflow.__name__}__{step_name}"
                workflow._tools.append((tool_name, make_tool(step_cls)))

        return Step

    @classmethod
    def get_next(cls, step_instance: object) -> object | None:
        """Get the next step.

        Calls the step's next() method with injected step classes and output.
        """
        if not hasattr(step_instance, "next"):
            return None

        next_method = step_instance.next
        sig = inspect.signature(next_method)

        # Inject step classes and result by parameter name
        kwargs = {}
        for param_name in sig.parameters:
            if param_name == "result":
                kwargs["result"] = getattr(step_instance, "output", None)
            elif param_name in cls._steps:
                kwargs[param_name] = cls._steps[param_name]
            else:
                step_name = type(step_instance).__name__
                msg = f"Unknown step '{param_name}' in {step_name}.next()"
                raise TypeError(msg)

        # next() returns a step class, we instantiate it with result from current step
        next_step_class = next_method(**kwargs)
        if next_step_class is None:
            return None
        return next_step_class(result=getattr(step_instance, "output", None))

    @classmethod
    def _write_step_file(
        cls, execution_id: str, step_instance: object, output: dict | None = None
    ):
        """Write step prompt and output to a file."""
        workflow_dir = cls._output_dirs.get(execution_id)
        if workflow_dir is None:
            return

        step_name = type(step_instance).__name__
        step_id = step_instance.id
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        filename = f"step-{step_name}-{timestamp}-{step_id}.json"

        data = {
            "step_name": step_name,
            "step_id": step_id,
            "prompt": step_instance.render(),
        }
        if output is not None:
            data["output"] = output

        (workflow_dir / filename).write_text(json.dumps(data, indent=2))

    @classmethod
    def start_workflow(cls) -> dict:
        """Start a new workflow execution and return the first step info."""
        if cls._start_step is None:
            raise RuntimeError(f"Workflow '{cls.__name__}' has no start step defined")

        execution_id = str(uuid.uuid4())

        # Create output directory for this workflow execution
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        dir_name = f"workflow-{cls.__name__}-{timestamp}-{execution_id}"
        workflow_dir = cls.output_dir / dir_name
        workflow_dir.mkdir(parents=True, exist_ok=True)
        cls._output_dirs[execution_id] = workflow_dir

        # Create new instance of start step
        step_class = cls._steps[cls._start_step]
        step_instance = step_class(execution_id=execution_id)

        cls._active_executions[execution_id] = step_instance

        # Write step file
        cls._write_step_file(execution_id, step_instance)

        first_tool_name = f"{cls.__name__}__{cls._start_step}"

        return {
            "execution_id": execution_id,
            "step_id": step_instance.id,
            "next_tool": first_tool_name,
            "next_step": cls._start_step,
            "prompt": step_instance.render(),
        }

    @classmethod
    def complete_step(
        cls, execution_id: str, step_id: str, status: str, output: dict[str, Any]
    ) -> dict:
        """Complete a step and return the next step's prompt."""
        if execution_id not in cls._active_executions:
            raise ValueError(f"Unknown execution_id: {execution_id}")

        current_step = cls._active_executions[execution_id]

        if current_step.id != step_id:
            msg = f"Step ID mismatch: expected {current_step.id}, got {step_id}"
            raise ValueError(msg)

        # Handle ERROR status - replay the current step
        if status == "ERROR":
            step_name = type(current_step).__name__
            step_uuid, attempt = current_step._parse_step_id(current_step.id)

            # Check if max retries exceeded
            if attempt >= current_step.max_retries:
                # Clean up
                del cls._active_executions[execution_id]
                if execution_id in cls._output_dirs:
                    del cls._output_dirs[execution_id]
                max_retries = current_step.max_retries
                return {
                    "execution_id": execution_id,
                    "status": "failed",
                    "message": f"Max retries ({max_retries}) exceeded for {step_name}",
                }

            # Increment attempt in step ID
            current_step.id = current_step._make_step_id(step_uuid, attempt + 1)

            cls._write_step_file(execution_id, current_step)
            return {
                "execution_id": execution_id,
                "step_id": current_step.id,
                "status": "retry",
                "attempt": attempt + 1,
                "max_retries": current_step.max_retries,
                "next_step": step_name,
                "prompt": current_step.render(),
            }

        # Validate output against result_schema
        if current_step.result_schema:
            try:
                validate(instance=output, schema=current_step.result_schema)
            except ValidationError as e:
                return {
                    "execution_id": execution_id,
                    "status": "validation_error",
                    "message": e.message,
                    "step": type(current_step).__name__,
                }

        # Store the output on the step
        current_step.output = output

        # Process artifacts from output
        if "artifacts" in output:
            workflow_dir = cls._output_dirs.get(execution_id)
            if workflow_dir is not None:
                if execution_id not in cls._artifacts:
                    cls._artifacts[execution_id] = []
                step_name = type(current_step).__name__
                for artifact_data in output["artifacts"]:
                    artifact = create_artifact(
                        artifact_data, workflow_dir, step_name, current_step.id
                    )
                    cls._artifacts[execution_id].append(artifact)

        # Write completed step file with output
        cls._write_step_file(execution_id, current_step, output)

        # Get next step and set execution_id
        next_step = cls.get_next(current_step)

        if next_step is not None:
            next_step.execution_id = execution_id

        if next_step is None:
            # Workflow complete
            del cls._active_executions[execution_id]
            if execution_id in cls._output_dirs:
                del cls._output_dirs[execution_id]
            return {
                "execution_id": execution_id,
                "status": "complete",
                "message": "Workflow completed successfully",
            }

        # Update active execution to next step
        cls._active_executions[execution_id] = next_step

        # Write next step file
        cls._write_step_file(execution_id, next_step)

        return {
            "execution_id": execution_id,
            "step_id": next_step.id,
            "next_step": type(next_step).__name__,
            "prompt": next_step.render(),
        }

    @classmethod
    def list_artifacts(cls, execution_id: str) -> dict:
        """List all artifacts for a workflow execution."""
        if execution_id not in cls._artifacts:
            return {"execution_id": execution_id, "artifacts": []}

        artifacts = cls._artifacts[execution_id]
        return {
            "execution_id": execution_id,
            "artifacts": [
                {"id": a.id, "name": a.name, "type": a.artifact_type} for a in artifacts
            ],
        }

    @classmethod
    def get_artifact(cls, execution_id: str, artifact_id: str) -> str:
        """Get the content of a specific artifact."""
        if execution_id not in cls._artifacts:
            raise ValueError(f"Unknown execution_id: {execution_id}")

        for artifact in cls._artifacts[execution_id]:
            if artifact.id == artifact_id:
                return artifact.read()

        raise ValueError(f"Unknown artifact_id: {artifact_id}")

    @classmethod
    def register(cls, mcp: "FastMCP"):
        """Register this workflow's tools and resources with a FastMCP instance."""
        # Register workflow-level tools
        start_tool_name = f"{cls.__name__}__start_workflow"
        mcp.tool(name=start_tool_name)(cls.start_workflow)

        complete_tool_name = f"{cls.__name__}__complete_step"
        mcp.tool(name=complete_tool_name)(cls.complete_step)

        # Register step tools
        for tool_name, bound_method in cls._tools:
            mcp.tool(name=tool_name)(bound_method)

        # Register artifact tools
        mcp.tool(name=f"{cls.__name__}__list_artifacts")(cls.list_artifacts)
        mcp.tool(name=f"{cls.__name__}__get_artifact")(cls.get_artifact)

        # Register workflow prompt
        def workflow_prompt() -> str:
            """Start a new workflow execution."""
            result = cls.start_workflow()
            return result["prompt"]

        mcp.prompt(name=cls.__name__)(workflow_prompt)
