from pathlib import Path

from fastmcp import FastMCP
from pydantic import BaseModel, Field

from artificer_workflows import Workflow

mcp = FastMCP(name="Artificer Workflows Development MCP", version="0.1.0")


# Shared model used by multiple steps
class ArtifactMetadata(BaseModel):
    """Metadata for artifacts created during workflow execution."""

    name: str = Field(description="Artifact filename (e.g., 'requirements.md')")
    path: str = Field(description="Path to the artifact")
    description: str | None = Field(
        default=None, description="Optional description of the artifact"
    )
    type: str | None = Field(
        default=None, description="Optional type hint (e.g., 'markdown', 'json')"
    )


# Workflow definition
class AddFeature(Workflow):
    templates_dir = Path(__file__).parent / "templates" / "add_feature"

    # Track artifacts across workflow execution
    _execution_artifacts: dict[str, list[dict]] = {}

    @classmethod
    def start_workflow(cls) -> dict:
        """Start workflow and initialize artifact tracking."""
        result = super().start_workflow()
        cls._execution_artifacts[result["workflow_id"]] = []
        return result

    @classmethod
    def complete_step(
        cls, workflow_id: str, step_id: str, status: str, output: dict
    ) -> dict:
        """Complete step and clean up artifacts when done."""
        result = super().complete_step(workflow_id, step_id, status, output)

        # Clean up artifacts when workflow completes or fails
        if result.get("status") in ("complete", "failed"):
            cls._execution_artifacts.pop(workflow_id, None)

        return result


class CollectRequirementsStep(AddFeature.Step, start=True):
    class OutputModel(BaseModel):
        summary: str = Field(description="Brief summary of collected requirements")
        artifact: ArtifactMetadata = Field(
            description="Requirements document artifact (e.g., requirements.md)"
        )

    def start(self, previous_result=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-CollectRequirementsStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("collect_requirements.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(self, output: OutputModel) -> type["ArchitectureReviewStep"]:
        artifact = output.artifact.model_dump()
        AddFeature._execution_artifacts[self.workflow_id].append(artifact)
        return ArchitectureReviewStep


class ArchitectureReviewStep(AddFeature.Step):
    class OutputModel(BaseModel):
        architecture_analysis: str = Field(
            description="Analysis of where the feature fits in current architecture"
        )
        simplifications_needed: list[str] = Field(
            description="List of things to simplify or remove before adding the feature"
        )
        refactoring_required: list[str] = Field(
            description="List of refactorings needed to support the feature"
        )
        architecture_changes: str = Field(
            description="Description of how the architecture will change"
        )
        artifact: ArtifactMetadata = Field(
            description="Updated ARCHITECTURE.md document"
        )

    def start(self, previous_result:CollectRequirementsStep.OutputModel=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-ArchitectureReviewStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("review_architecture.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(self, output: OutputModel) -> type["CreatePlanStep"]:
        artifact = output.artifact.model_dump()
        AddFeature._execution_artifacts[self.workflow_id].append(artifact)
        return CreatePlanStep


class CreatePlanStep(AddFeature.Step):
    class OutputModel(BaseModel):
        summary: str = Field(description="Brief summary of the implementation plan")
        artifact: ArtifactMetadata = Field(
            description="Implementation plan artifact (e.g., plan.md)"
        )

    def start(self, previous_result:ArchitectureReviewStep.OutputModel=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-CreatePlanStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("create_plan.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(self, output: OutputModel) -> type["ImplementFeatureStep"]:
        artifact = output.artifact.model_dump()
        AddFeature._execution_artifacts[self.workflow_id].append(artifact)
        return ImplementFeatureStep


class ImplementFeatureStep(AddFeature.Step):
    class OutputModel(BaseModel):
        summary: str = Field(description="Brief summary of what was implemented")
        artifacts: list[ArtifactMetadata] = Field(
            min_length=1,
            description="List of files/artifacts created during implementation",
        )

    def start(self, previous_result:ArchitectureReviewStep.OutputModel=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-ImplementFeatureStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("implement_feature.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(self, output: OutputModel) -> type["TestFeatureStep"]:
        for artifact in output.artifacts:
            artifact_dict = artifact.model_dump()
            AddFeature._execution_artifacts[self.workflow_id].append(artifact_dict)
        return TestFeatureStep


class TestFeatureStep(AddFeature.Step):
    class OutputModel(BaseModel):
        passed: bool = Field(description="Whether all tests passed")
        summary: str = Field(description="Brief summary of test results")
        artifact: ArtifactMetadata | None = Field(
            default=None,
            description="Optional test results artifact (e.g., test_output.txt)",
        )

    def start(self, previous_result:ImplementFeatureStep.OutputModel=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-TestFeatureStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("test_feature.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(self, output: OutputModel) -> type["ReviewFeatureStep"]:
        if output.artifact:
            artifact = output.artifact.model_dump()
            AddFeature._execution_artifacts[self.workflow_id].append(artifact)
        return ReviewFeatureStep


class ReviewFeatureStep(AddFeature.Step):
    class OutputModel(BaseModel):
        needs_revision: bool = Field(
            description="Whether the feature needs revision based on review"
        )
        summary: str = Field(description="Brief summary of review findings")
        artifact: ArtifactMetadata = Field(
            description="Code review artifact (e.g., review.md)"
        )

    def start(self, previous_result:TestFeatureStep.OutputModel=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-ReviewFeatureStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("review_feature.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(
        self, output: OutputModel
    ) -> type["CreatePlanStep"] | type["SummaryStep"]:
        artifact = output.artifact.model_dump()
        AddFeature._execution_artifacts[self.workflow_id].append(artifact)

        if output.needs_revision:
            return CreatePlanStep
        return SummaryStep


class SummaryStep(AddFeature.Step):
    class OutputModel(BaseModel):
        summary: str = Field(description="Brief summary of completed feature work")
        artifact: ArtifactMetadata = Field(
            description="Summary document artifact (e.g., summary.md)"
        )

    def start(self, previous_result:ReviewFeatureStep.OutputModel=None) -> str:
        artifacts = AddFeature._execution_artifacts.get(self.workflow_id, [])
        from artificer_workflows.store import workflow_store

        workflow = workflow_store.get_workflow(self.workflow_id)
        step_dir = f".artificer/workflow-AddFeature-{workflow.start_time}-{self.workflow_id}/step-SummaryStep-{self.start_time}-{self.step_id}"
        template = AddFeature._jinja_env.get_template("summary.md")
        return template.render(
            result=previous_result, artifacts=artifacts, step_dir=step_dir
        )

    def complete(self, output: OutputModel) -> None:
        artifact = output.artifact.model_dump()
        AddFeature._execution_artifacts[self.workflow_id].append(artifact)
        return None


AddFeature.register(mcp)


if __name__ == "__main__":
    mcp.run()
