from __future__ import annotations

import difflib
import os
import re
import tempfile
from pathlib import Path

from ai_orchestrator.schemas.tool_result import ToolResult
from ai_orchestrator.tools.shell_runner import ShellRunner


class GitRunner:
    _HUNK_HEADER = re.compile(
        r"^@@ -(?P<old_start>\d+)(?:,(?P<old_count>\d+))? \+(?P<new_start>\d+)(?:,(?P<new_count>\d+))? @@"
    )
    _IGNORED_PARTS = {".git", ".venv", "node_modules", "__pycache__", "artifacts", "data"}

    def __init__(self, shell_runner: ShellRunner) -> None:
        self.shell_runner = shell_runner

    async def apply_patch(
        self,
        *,
        patch: str,
        repo_path: Path,
        timeout_seconds: int | None = None,
    ) -> ToolResult:
        if not patch.strip():
            return ToolResult(success=True, stdout="No patch to apply.", command="apply patch")

        normalized_patch = self._normalize_patch(patch)
        if (repo_path / ".git").exists():
            git_result = await self._git_apply_patch(
                patch=normalized_patch,
                repo_path=repo_path,
                timeout_seconds=timeout_seconds,
            )
            if git_result.success:
                return git_result
            fallback = self._apply_patch_fallback(normalized_patch, repo_path)
            if fallback is not None:
                fallback.data["git_apply_stderr"] = git_result.stderr
                return fallback
            return git_result

        fallback = self._apply_patch_fallback(normalized_patch, repo_path)
        if fallback is not None:
            return fallback
        return ToolResult(
            success=False,
            stderr="Patch could not be applied by the internal patch engine.",
            command="apply patch",
        )

    async def git_apply_patch(
        self,
        *,
        patch: str,
        repo_path: Path,
        timeout_seconds: int | None = None,
    ) -> ToolResult:
        return await self.apply_patch(
            patch=patch,
            repo_path=repo_path,
            timeout_seconds=timeout_seconds,
        )

    async def _git_apply_patch(
        self,
        *,
        patch: str,
        repo_path: Path,
        timeout_seconds: int | None = None,
    ) -> ToolResult:
        file_descriptor, patch_path = tempfile.mkstemp(suffix=".diff", dir=repo_path)
        try:
            os.close(file_descriptor)
            Path(patch_path).write_text(patch, encoding="utf-8-sig", newline="\n")
            return await self.shell_runner.run_shell(
                command=f'git apply --whitespace=nowarn "{Path(patch_path).name}"',
                repo_path=repo_path,
                timeout_seconds=timeout_seconds,
                use_docker=False,
            )
        finally:
            Path(patch_path).unlink(missing_ok=True)

    async def diff_repo(
        self,
        *,
        repo_path: Path,
        base_path: Path | None = None,
        timeout_seconds: int | None = None,
    ) -> ToolResult:
        if base_path is not None:
            diff_text = self._directory_diff(base_path, repo_path)
            return ToolResult(success=True, stdout=diff_text, command="unified diff")
        if (repo_path / ".git").exists():
            return await self._git_diff(repo_path=repo_path, timeout_seconds=timeout_seconds)
        return ToolResult(success=True, stdout="", command="unified diff")

    async def git_diff(
        self,
        *,
        repo_path: Path,
        timeout_seconds: int | None = None,
    ) -> ToolResult:
        return await self.diff_repo(repo_path=repo_path, timeout_seconds=timeout_seconds)

    async def _git_diff(
        self,
        *,
        repo_path: Path,
        timeout_seconds: int | None = None,
    ) -> ToolResult:
        return await self.shell_runner.run_shell(
            command="git diff --no-ext-diff --binary -- .",
            repo_path=repo_path,
            timeout_seconds=timeout_seconds,
            use_docker=False,
        )

    @staticmethod
    def _normalize_patch(patch: str) -> str:
        if not patch.strip() or "diff --git " in patch:
            return patch

        lines = patch.splitlines()
        normalized: list[str] = []
        index = 0
        while index < len(lines):
            line = lines[index]
            if line.startswith("--- ") and index + 1 < len(lines) and lines[index + 1].startswith("+++ "):
                old_path = lines[index][4:].strip()
                new_path = lines[index + 1][4:].strip()
                normalized.append(f"diff --git {old_path} {new_path}")
                normalized.append(line)
                normalized.append(lines[index + 1])
                index += 2
                continue
            normalized.append(line)
            index += 1

        normalized_patch = "\n".join(normalized)
        if patch.endswith("\n"):
            normalized_patch += "\n"
        return normalized_patch

    def _apply_patch_fallback(self, patch: str, repo_path: Path) -> ToolResult | None:
        sections = self._parse_patch_sections(patch)
        if not sections:
            return None

        for section in sections:
            old_relative = self._strip_patch_prefix(str(section["old_path"]))
            new_relative = self._strip_patch_prefix(str(section["new_path"]))
            is_delete = str(section["new_path"]).strip() == "/dev/null"
            is_add = str(section["old_path"]).strip() == "/dev/null"
            target_path = repo_path / (old_relative if is_delete else new_relative)
            if target_path.exists():
                original_lines = target_path.read_text(encoding="utf-8").splitlines(keepends=True)
            elif is_add:
                original_lines = []
            else:
                return None
            updated_lines = list(original_lines)
            line_offset = 0
            for hunk in section["hunks"]:
                old_lines = self._materialize_hunk_lines(hunk["lines"], include_prefixes={" ", "-"})
                new_lines = self._materialize_hunk_lines(hunk["lines"], include_prefixes={" ", "+"})
                start_hint = max(0, hunk["old_start"] - 1 + line_offset)
                position = self._locate_hunk(updated_lines, old_lines, start_hint=start_hint)
                if position is None:
                    return None
                end = position + len(old_lines)
                updated_lines[position:end] = new_lines
                line_offset += len(new_lines) - len(old_lines)
            if is_delete:
                target_path.unlink(missing_ok=True)
                continue
            if updated_lines != original_lines:
                target_path.parent.mkdir(parents=True, exist_ok=True)
                target_path.write_text("".join(updated_lines), encoding="utf-8", newline="\n")

        return ToolResult(
            success=True,
            stdout="Applied patch with fallback patcher.",
            command="internal apply patch",
            data={"fallback": True},
        )

    @staticmethod
    def _parse_patch_sections(patch: str) -> list[dict[str, object]]:
        lines = patch.splitlines()
        sections: list[dict[str, object]] = []
        current: dict[str, object] | None = None
        index = 0
        while index < len(lines):
            line = lines[index]
            if line.startswith("diff --git "):
                index += 1
                continue
            if line.startswith("--- "):
                if index + 1 >= len(lines) or not lines[index + 1].startswith("+++ "):
                    return []
                current = {
                    "old_path": lines[index][4:].strip(),
                    "new_path": lines[index + 1][4:].strip(),
                    "hunks": [],
                }
                sections.append(current)
                index += 2
                continue
            if line.startswith("@@ "):
                if current is None:
                    return []
                match = GitRunner._HUNK_HEADER.match(line)
                if not match:
                    return []
                hunk_lines: list[str] = [line]
                index += 1
                while index < len(lines):
                    candidate = lines[index]
                    if candidate.startswith("diff --git ") or candidate.startswith("--- ") or candidate.startswith("@@ "):
                        break
                    hunk_lines.append(candidate)
                    index += 1
                current["hunks"].append(
                    {
                        "old_start": int(match.group("old_start")),
                        "new_start": int(match.group("new_start")),
                        "lines": hunk_lines,
                    }
                )
                continue
            index += 1
        return sections

    @staticmethod
    def _materialize_hunk_lines(hunk_lines: list[str], include_prefixes: set[str]) -> list[str]:
        materialized: list[str] = []
        for line in hunk_lines[1:]:
            if not line:
                materialized.append("\n")
                continue
            if line.startswith("\\"):
                continue
            prefix = line[0]
            if prefix in include_prefixes:
                materialized.append(f"{line[1:]}\n")
        return materialized

    @staticmethod
    def _locate_hunk(file_lines: list[str], old_lines: list[str], *, start_hint: int) -> int | None:
        if not old_lines:
            return min(max(start_hint, 0), len(file_lines))

        max_start = len(file_lines) - len(old_lines)
        if max_start < 0:
            return None

        for delta in range(0, 8):
            for candidate in {start_hint - delta, start_hint + delta}:
                if 0 <= candidate <= max_start and file_lines[candidate : candidate + len(old_lines)] == old_lines:
                    return candidate

        for candidate in range(max_start + 1):
            if file_lines[candidate : candidate + len(old_lines)] == old_lines:
                return candidate
        return None

    def _directory_diff(self, before_root: Path, after_root: Path) -> str:
        before_files = self._text_files(before_root)
        after_files = self._text_files(after_root)
        relative_paths = sorted(set(before_files) | set(after_files))
        chunks: list[str] = []

        for relative_path in relative_paths:
            before_lines = before_files.get(relative_path, [])
            after_lines = after_files.get(relative_path, [])
            if before_lines == after_lines:
                continue
            diff = difflib.unified_diff(
                before_lines,
                after_lines,
                fromfile=f"a/{relative_path}",
                tofile=f"b/{relative_path}",
                lineterm="",
            )
            chunk = "\n".join(diff)
            if chunk:
                chunks.append(chunk)
        return "\n".join(chunks) + ("\n" if chunks else "")

    def _text_files(self, root: Path) -> dict[str, list[str]]:
        files: dict[str, list[str]] = {}
        if not root.exists():
            return files
        for file_path in root.rglob("*"):
            if not file_path.is_file():
                continue
            relative = file_path.relative_to(root)
            if any(part in self._IGNORED_PARTS for part in relative.parts):
                continue
            try:
                content = file_path.read_text(encoding="utf-8")
            except (OSError, UnicodeDecodeError):
                continue
            files[relative.as_posix()] = content.splitlines(keepends=True)
        return files

    @staticmethod
    def _strip_patch_prefix(path: str) -> Path:
        cleaned = path.strip()
        if cleaned in {"", "/dev/null"}:
            return Path(cleaned)
        if cleaned.startswith(("a/", "b/")):
            cleaned = cleaned[2:]
        return Path(cleaned.replace("/", os.sep))
