from __future__ import annotations

import asyncio
import shutil
from pathlib import Path

from ai_orchestrator.config.settings import Settings
from ai_orchestrator.schemas.tool_result import ToolResult
from ai_orchestrator.schemas.workflow_models import WorkspaceSnapshotRecord
from ai_orchestrator.tools.git_runner import GitRunner
from ai_orchestrator.tools.shell_runner import ShellRunner


class WorkspaceManager:
    def __init__(
        self,
        *,
        settings: Settings,
        shell_runner: ShellRunner,
        git_runner: GitRunner,
    ) -> None:
        self.settings = settings
        self.shell_runner = shell_runner
        self.git_runner = git_runner

    async def create_snapshot(self, task_id: str, repo_path: Path) -> WorkspaceSnapshotRecord:
        workspace_path = self.settings.workspaces_dir / task_id
        if workspace_path.exists():
            await asyncio.to_thread(shutil.rmtree, workspace_path, True)

        if (repo_path / ".git").exists():
            head_check = await self.shell_runner.run_shell(
                command="git rev-parse --verify HEAD",
                repo_path=repo_path,
                use_docker=False,
            )
            if head_check.success:
                result = await self.shell_runner.run_shell(
                    command=f'git worktree add --detach "{workspace_path}" HEAD',
                    repo_path=repo_path,
                    use_docker=False,
                )
                if result.success:
                    return WorkspaceSnapshotRecord(
                        task_id=task_id,
                        source_repo_path=str(repo_path),
                        workspace_path=str(workspace_path),
                        strategy="git_worktree",
                    )

        await asyncio.to_thread(
            shutil.copytree,
            repo_path,
            workspace_path,
            dirs_exist_ok=True,
            ignore=shutil.ignore_patterns(
                "__pycache__",
                ".pytest_cache",
                "pytest-cache-files-*",
                "artifacts",
                "data",
                ".git",
            ),
        )
        return WorkspaceSnapshotRecord(
            task_id=task_id,
            source_repo_path=str(repo_path),
            workspace_path=str(workspace_path),
            strategy="copy",
        )

    async def restore_snapshot(self, snapshot: WorkspaceSnapshotRecord) -> None:
        await self.commit_snapshot(snapshot, Path(snapshot.source_repo_path))

    async def commit_snapshot(self, snapshot: WorkspaceSnapshotRecord, destination_repo: Path) -> ToolResult:
        workspace_path = Path(snapshot.workspace_path)
        if snapshot.strategy == "git_worktree":
            diff_result = await self.git_runner.diff_repo(repo_path=workspace_path)
            if not diff_result.success or not diff_result.stdout.strip():
                await self.discard_snapshot(snapshot)
                return diff_result
            apply_result = await self.git_runner.apply_patch(
                patch=diff_result.stdout,
                repo_path=destination_repo,
            )
            if apply_result.success:
                await self._copy_untracked_files(workspace_path, destination_repo)
            await self.discard_snapshot(snapshot)
            apply_result.data["diff"] = diff_result.stdout
            return apply_result

        diff_result = await self.git_runner.diff_repo(repo_path=workspace_path, base_path=destination_repo)
        diff_text = diff_result.stdout if diff_result.success else ""
        await asyncio.to_thread(self._sync_tree, workspace_path, destination_repo)
        await self.discard_snapshot(snapshot)
        return ToolResult(success=True, stdout=diff_text, data={"diff": diff_text})

    async def discard_snapshot(self, snapshot: WorkspaceSnapshotRecord) -> None:
        workspace_path = Path(snapshot.workspace_path)
        if snapshot.strategy == "git_worktree":
            await self.shell_runner.run_shell(
                command=f'git worktree remove --force "{workspace_path}"',
                repo_path=Path(snapshot.source_repo_path),
                use_docker=False,
            )
            return
        await asyncio.to_thread(shutil.rmtree, workspace_path, True)

    async def _copy_untracked_files(self, workspace_path: Path, destination_repo: Path) -> None:
        result = await self.shell_runner.run_shell(
            command="git ls-files --others --exclude-standard",
            repo_path=workspace_path,
            use_docker=False,
        )
        if not result.success:
            return
        for line in result.stdout.splitlines():
            relative_path = Path(line.strip())
            if not line.strip():
                continue
            source = workspace_path / relative_path
            destination = destination_repo / relative_path
            if source.is_dir():
                continue
            destination.parent.mkdir(parents=True, exist_ok=True)
            await asyncio.to_thread(shutil.copy2, source, destination)

    def _sync_tree(self, source_root: Path, destination_root: Path) -> str:
        changed_files: list[str] = []
        for path in source_root.rglob("*"):
            if not path.is_file():
                continue
            relative = path.relative_to(source_root)
            if "__pycache__" in relative.parts:
                continue
            destination = destination_root / relative
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(path, destination)
            changed_files.append(relative.as_posix())
        return "\n".join(changed_files)
