"""Merge executor — performs git merge with distributed locking and conflict detection.

Also provides ``execute_merge_task`` which wraps the core merge with the full
post-merge lifecycle: test, push, worktree cleanup, and task status updates.
"""

from __future__ import annotations

import asyncio
import logging
import os
from dataclasses import dataclass, field
from enum import StrEnum
from pathlib import Path

import asyncpg
from redis.asyncio import Redis

from loom.bus.events import EventType
from loom.bus.locks import acquire_merge_lock, release_merge_lock
from loom.bus.publisher import publish_event
from loom.graph import cache, store

log = logging.getLogger(__name__)

# Maximum characters of pytest output to store in task output / failure reason.
_MAX_TEST_OUTPUT_CHARS = 8000


# Merge phase names for checkpointing
class MergePhase(StrEnum):
    LOCK_ACQUIRED = "lock_acquired"
    MERGED = "merged"
    TESTS_RUNNING = "tests_running"
    TESTS_PASSED = "tests_passed"
    PUSHED = "pushed"
    CLEANED_UP = "cleaned_up"


class MergeStatus(StrEnum):
    SUCCESS = "success"
    CONFLICT = "conflict"
    LOCK_FAILED = "lock_failed"
    GIT_ERROR = "git_error"
    TEST_FAILED = "test_failed"


@dataclass
class ConflictInfo:
    """Details about a merge conflict."""

    conflicting_files: list[str] = field(default_factory=list)
    auto_resolvable: bool = False
    details: str = ""


@dataclass
class MergeResult:
    """Result of a merge operation."""

    status: MergeStatus
    branch_name: str
    message: str = ""
    conflict: ConflictInfo | None = None
    commit_sha: str = ""
    test_output: str = ""

    def to_dict(self) -> dict:
        result: dict = {
            "status": self.status.value,
            "branch_name": self.branch_name,
            "message": self.message,
        }
        if self.commit_sha:
            result["commit_sha"] = self.commit_sha
        if self.conflict:
            result["conflict"] = {
                "conflicting_files": self.conflict.conflicting_files,
                "auto_resolvable": self.conflict.auto_resolvable,
                "details": self.conflict.details,
            }
        if self.test_output:
            result["test_output"] = self.test_output
        return result


async def _run_git(
    *args: str, cwd: str | Path, check: bool = True
) -> tuple[int, str, str]:
    """Run a git command asynchronously.

    Returns (returncode, stdout, stderr).
    Raises RuntimeError if check=True and returncode != 0.
    """
    cmd = ["git"] + list(args)
    log.debug("Running: %s (cwd=%s)", " ".join(cmd), cwd)

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        cwd=str(cwd),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env={**os.environ, "GIT_TERMINAL_PROMPT": "0"},
    )
    stdout_bytes, stderr_bytes = await proc.communicate()
    stdout = stdout_bytes.decode().strip()
    stderr = stderr_bytes.decode().strip()
    returncode = proc.returncode or 0

    if check and returncode != 0:
        raise RuntimeError(
            f"git {args[0]} failed (rc={returncode}): {stderr or stdout}"
        )

    log.debug("git %s rc=%d stdout=%r stderr=%r", args[0], returncode, stdout, stderr)
    return returncode, stdout, stderr


async def _get_conflicting_files(project_dir: str | Path) -> list[str]:
    """After a failed merge, find the conflicting files."""
    rc, stdout, _ = await _run_git(
        "diff", "--name-only", "--diff-filter=U",
        cwd=project_dir, check=False,
    )
    if rc != 0 or not stdout:
        return []
    return [f.strip() for f in stdout.splitlines() if f.strip()]


async def _is_simple_conflict(
    conflicting_files: list[str], branch_name: str, project_dir: str | Path
) -> bool:
    """Determine if a conflict is simple (non-overlapping changes).

    A simple conflict is one where:
    - There are conflicting files, but each file was only modified on one side.
      This can happen when git marks a delete/modify as a conflict.
    - We check by seeing if any file has actual merge conflict markers.

    Returns True if auto-resolvable, False otherwise.
    """
    if not conflicting_files:
        return False

    for filepath in conflicting_files:
        full_path = Path(project_dir) / filepath
        if not full_path.exists():
            # File was deleted on one side -- that is resolvable by accepting theirs
            continue
        try:
            content = full_path.read_text(errors="replace")
            # If we find conflict markers, it is a real content conflict
            if "<<<<<<< HEAD" in content and "=======" in content and ">>>>>>>" in content:
                return False
        except OSError:
            return False

    return True


async def _auto_resolve_conflicts(
    conflicting_files: list[str], branch_name: str, project_dir: str | Path
) -> bool:
    """Auto-resolve simple conflicts by accepting the branch version.

    For delete/modify conflicts or other simple cases, accept the incoming
    (branch) version of each file.

    Returns True if all conflicts were resolved successfully.
    """
    try:
        for filepath in conflicting_files:
            # Accept the branch (theirs) version
            await _run_git("checkout", "--theirs", "--", filepath, cwd=project_dir)
            await _run_git("add", filepath, cwd=project_dir)
        return True
    except RuntimeError:
        log.warning("Auto-resolve failed for %s", filepath, exc_info=True)
        return False


async def _abort_merge(project_dir: str | Path) -> None:
    """Abort an in-progress merge, best-effort."""
    try:
        await _run_git("merge", "--abort", cwd=project_dir, check=False)
    except Exception:
        log.warning("Failed to abort merge", exc_info=True)


async def execute_merge(
    redis_client: Redis,
    branch_name: str,
    project_dir: str | Path,
    target_branch: str = "main",
    lock_ttl: int = 300,
    auto_resolve: bool = True,
    run_tests_cmd: str | None = None,
) -> MergeResult:
    """Execute a full merge workflow with distributed locking.

    Steps:
    1. Acquire distributed lock on target branch
    2. Checkout target branch and pull latest
    3. Merge the feature branch
    4. Optionally run tests
    5. Handle conflicts (auto-resolve if simple)
    6. Always release lock in finally block

    Args:
        redis_client: Async Redis connection for distributed locking.
        branch_name: The feature branch to merge.
        project_dir: Path to the git repository.
        target_branch: Branch to merge into (default: "main").
        lock_ttl: Lock timeout in seconds (default: 300).
        auto_resolve: Whether to attempt auto-resolution of simple conflicts.
        run_tests_cmd: Optional shell command to run tests after merge.

    Returns:
        MergeResult with status and details.
    """
    project_dir = Path(project_dir)
    lock_token: str | None = None

    try:
        # Step 1: Acquire the distributed merge lock
        log.info("Acquiring merge lock for branch %s", branch_name)
        lock_token = await acquire_merge_lock(redis_client, ttl_seconds=lock_ttl)
        if lock_token is None:
            log.warning("Could not acquire merge lock — another merge in progress")
            return MergeResult(
                status=MergeStatus.LOCK_FAILED,
                branch_name=branch_name,
                message="Could not acquire merge lock. Another merge is in progress.",
            )

        # Step 2: Checkout target branch and pull latest
        log.info("Checking out %s and pulling latest", target_branch)
        await _run_git("checkout", target_branch, cwd=project_dir)
        # Pull but don't fail if remote is not configured
        rc, _, stderr = await _run_git(
            "pull", "--ff-only", cwd=project_dir, check=False
        )
        if rc != 0 and "no tracking information" not in stderr.lower():
            # Only log as warning; local-only repos won't have a remote
            log.debug("git pull returned rc=%d: %s", rc, stderr)

        # Step 3: Merge the feature branch
        log.info("Merging branch %s into %s", branch_name, target_branch)
        rc, stdout, stderr = await _run_git(
            "merge", "--no-edit", branch_name, cwd=project_dir, check=False
        )

        if rc != 0:
            # Merge failed — check for conflicts
            conflicting_files = await _get_conflicting_files(project_dir)

            if conflicting_files:
                log.info(
                    "Merge conflict detected in %d file(s): %s",
                    len(conflicting_files),
                    conflicting_files,
                )

                conflict_info = ConflictInfo(
                    conflicting_files=conflicting_files,
                    auto_resolvable=False,
                    details=stderr or stdout,
                )

                # Check if conflicts are auto-resolvable
                if auto_resolve:
                    is_simple = await _is_simple_conflict(
                        conflicting_files, branch_name, project_dir
                    )
                    if is_simple:
                        resolved = await _auto_resolve_conflicts(
                            conflicting_files, branch_name, project_dir
                        )
                        if resolved:
                            conflict_info.auto_resolvable = True
                            # Complete the merge commit
                            await _run_git(
                                "commit", "--no-edit", "-m",
                                f"Merge branch '{branch_name}' (auto-resolved)",
                                cwd=project_dir,
                            )
                            # Get the commit SHA
                            _, sha, _ = await _run_git(
                                "rev-parse", "HEAD", cwd=project_dir
                            )
                            log.info(
                                "Auto-resolved conflicts and committed: %s", sha
                            )
                            return MergeResult(
                                status=MergeStatus.SUCCESS,
                                branch_name=branch_name,
                                message=f"Merged with auto-resolved conflicts in {len(conflicting_files)} file(s)",
                                conflict=conflict_info,
                                commit_sha=sha,
                            )

                # Complex conflict — abort and report
                await _abort_merge(project_dir)
                return MergeResult(
                    status=MergeStatus.CONFLICT,
                    branch_name=branch_name,
                    message=f"Merge conflict in {len(conflicting_files)} file(s) that could not be auto-resolved",
                    conflict=conflict_info,
                )
            else:
                # Some other git error (not a conflict)
                return MergeResult(
                    status=MergeStatus.GIT_ERROR,
                    branch_name=branch_name,
                    message=f"git merge failed: {stderr or stdout}",
                )

        # Step 4: Optionally run tests
        if run_tests_cmd:
            log.info("Running tests: %s", run_tests_cmd)
            proc = await asyncio.create_subprocess_shell(
                run_tests_cmd,
                cwd=str(project_dir),
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            test_stdout, test_stderr = await proc.communicate()
            if proc.returncode != 0:
                # Tests failed — abort the merge by resetting
                log.warning("Tests failed after merge, resetting")
                await _run_git("reset", "--hard", f"HEAD~1", cwd=project_dir)
                return MergeResult(
                    status=MergeStatus.GIT_ERROR,
                    branch_name=branch_name,
                    message=f"Tests failed after merge: {test_stderr.decode().strip()}",
                )

        # Step 5: Get the merge commit SHA
        _, sha, _ = await _run_git("rev-parse", "HEAD", cwd=project_dir)

        log.info("Merge of %s into %s succeeded: %s", branch_name, target_branch, sha)
        return MergeResult(
            status=MergeStatus.SUCCESS,
            branch_name=branch_name,
            message=f"Successfully merged {branch_name} into {target_branch}",
            commit_sha=sha,
        )

    except RuntimeError as exc:
        log.error("Git error during merge of %s: %s", branch_name, exc)
        # Best-effort abort if merge was in progress
        await _abort_merge(project_dir)
        return MergeResult(
            status=MergeStatus.GIT_ERROR,
            branch_name=branch_name,
            message=str(exc),
        )

    finally:
        # Always release the lock
        if lock_token is not None:
            try:
                released = await release_merge_lock(redis_client, lock_token)
                if released:
                    log.info("Released merge lock (token=%s...)", lock_token[:8])
                else:
                    log.warning("Lock release returned False (expired or stolen?)")
            except Exception:
                log.error("Failed to release merge lock", exc_info=True)


# ---------------------------------------------------------------------------
# Post-merge lifecycle helpers
# ---------------------------------------------------------------------------


def _truncate_output(text: str) -> str:
    """Truncate test output to the last ``_MAX_TEST_OUTPUT_CHARS`` characters."""
    if len(text) <= _MAX_TEST_OUTPUT_CHARS:
        return text
    return "... (truncated)\n" + text[-_MAX_TEST_OUTPUT_CHARS:]


async def _run_tests(
    project_dir: str | Path,
    test_cmd: str = "uv run pytest tests/ -v -k 'not integration'",
) -> tuple[bool, str]:
    """Run the test suite after merge.

    Returns (passed: bool, output: str).  Output is truncated to
    ``_MAX_TEST_OUTPUT_CHARS``.
    """
    log.info("Running post-merge tests: %s", test_cmd)
    proc = await asyncio.create_subprocess_shell(
        test_cmd,
        cwd=str(project_dir),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,  # combine for a single stream
    )
    stdout_bytes, _ = await proc.communicate()
    output = stdout_bytes.decode(errors="replace").strip()
    output = _truncate_output(output)
    passed = proc.returncode == 0
    if passed:
        log.info("Post-merge tests passed")
    else:
        log.warning("Post-merge tests FAILED (rc=%d)", proc.returncode)
    return passed, output


async def _push_to_origin(project_dir: str | Path) -> None:
    """Push the current branch (main) to origin. Best-effort."""
    rc, _, stderr = await _run_git("push", "origin", "HEAD", cwd=project_dir, check=False)
    if rc == 0:
        log.info("Pushed to origin")
    else:
        # Local-only repos or missing remote — not fatal.
        log.debug("git push returned rc=%d: %s (non-fatal)", rc, stderr)


async def _cleanup_worktree(
    project_dir: str | Path,
    worktree_path: str,
    branch_name: str,
) -> None:
    """Remove the worktree and delete the feature branch.  Best-effort for
    each step — already-cleaned state should not cause a hard failure.

    Order matters: worktree remove *before* branch delete.
    """
    # 1. Remove worktree (ignore if already gone)
    rc, _, stderr = await _run_git(
        "worktree", "remove", "--force", worktree_path,
        cwd=project_dir, check=False,
    )
    if rc != 0:
        log.debug("git worktree remove rc=%d: %s (ignored)", rc, stderr)

    # 2. Delete the feature branch (ignore if already gone)
    rc, _, stderr = await _run_git(
        "branch", "-D", branch_name,
        cwd=project_dir, check=False,
    )
    if rc != 0:
        log.debug("git branch -D rc=%d: %s (ignored)", rc, stderr)


async def _heartbeat_loop(pool: asyncpg.Pool, task_id: str, interval: int = 120):
    """Extend claim TTL periodically while tests run.

    Runs as a background task, cancelled when tests complete.
    """
    from datetime import datetime, timedelta, timezone

    while True:
        await asyncio.sleep(interval)
        try:
            expires_at = datetime.now(timezone.utc) + timedelta(seconds=600)
            await store.set_claim_expiry(pool, task_id, expires_at)
            log.debug("Heartbeat: extended TTL for task %s", task_id)
        except Exception:
            log.warning("Heartbeat failed for task %s", task_id, exc_info=True)


async def _validate_branch(
    branch_name: str,
    project_dir: str | Path,
    target_branch: str = "main",
) -> tuple[bool, str]:
    """Pre-merge branch validation.

    Checks:
    1. Branch exists locally
    2. Branch has commits ahead of target

    Returns (valid, message).
    """
    # Check branch exists
    rc, _, stderr = await _run_git(
        "rev-parse", "--verify", branch_name,
        cwd=project_dir, check=False,
    )
    if rc != 0:
        return False, f"Branch '{branch_name}' does not exist"

    # Check branch has new commits
    rc, stdout, _ = await _run_git(
        "log", f"{target_branch}..{branch_name}", "--oneline",
        cwd=project_dir, check=False,
    )
    if rc != 0:
        return False, f"Cannot compare {branch_name} with {target_branch}"
    if not stdout.strip():
        return False, f"Branch '{branch_name}' has no new commits ahead of {target_branch}"

    return True, f"Branch has {len(stdout.strip().splitlines())} commit(s) ahead"


async def _set_merge_phase(
    pool: asyncpg.Pool,
    task_id: str,
    phase: str,
) -> None:
    """Store the current merge phase in the task's context for crash recovery."""
    try:
        task = await store.get_task(pool, task_id)
        ctx = dict(task.context) if task.context else {}
        ctx["merge_phase"] = phase
        await store.update_task(pool, task_id, context=ctx)
    except Exception:
        log.warning(
            "Failed to checkpoint phase %s for task %s", phase, task_id, exc_info=True
        )


async def execute_merge_task(
    pool: asyncpg.Pool,
    redis_client: Redis,
    project_id: str,
    task_id: str,
    branch_name: str,
    project_dir: str | Path,
    worktree_path: str,
    target_branch: str = "main",
    test_cmd: str = "uv run pytest tests/ -v -k 'not integration'",
) -> MergeResult:
    """Full merge-task lifecycle: merge, test, push, cleanup, update task.

    Phases executed in order:

    1. **Merge** — delegates to :func:`execute_merge`.
    2. **Test**  — runs ``test_cmd``, captures output (truncated to 8 kB).
    3. **Success path** — push to origin, remove worktree, delete branch,
       mark task done via ``store.complete_task``, publish success event.
    4. **Failure path** — ``git reset --hard HEAD~1``, mark task failed via
       ``store.fail_task``, publish failure event.

    The Redis merge lock is acquired/released inside :func:`execute_merge` so
    it is always freed regardless of outcome.  Worktree cleanup runs in a
    ``finally`` block so it executes in both success and failure paths.

    Args:
        pool: asyncpg connection pool for Postgres writes.
        redis_client: async Redis connection (lock + events).
        project_id: Loom project UUID string.
        task_id: The merge task's loom-{hex} ID (must be in claimed state).
        branch_name: Feature branch to merge.
        project_dir: Path to the git repo root.
        worktree_path: Relative path to the worktree directory.
        target_branch: Branch to merge into (default ``"main"``).
        test_cmd: Shell command to run tests after merge.

    Returns:
        :class:`MergeResult` with full details.
    """
    project_dir = Path(project_dir)
    cleanup_needed = False

    try:
        # ----- Pre-merge: Branch validation -----
        valid, msg = await _validate_branch(branch_name, project_dir, target_branch)
        if not valid:
            reason = f"Branch validation failed: {msg}"
            failed_task = await store.fail_task(pool, task_id, reason)
            await cache.sync_task(redis_client, failed_task)
            await publish_event(
                redis_client, project_id, EventType.MERGE_FAILED,
                task_id=task_id,
                payload={"reason": reason, "branch_name": branch_name},
            )
            return MergeResult(
                status=MergeStatus.GIT_ERROR,
                branch_name=branch_name,
                message=reason,
            )

        # ----- Phase 1: Merge -----
        await publish_event(
            redis_client, project_id, EventType.MERGE_STARTED,
            task_id=task_id,
            payload={"branch_name": branch_name},
        )
        await _set_merge_phase(pool, task_id, MergePhase.LOCK_ACQUIRED)

        merge_result = await execute_merge(
            redis_client=redis_client,
            branch_name=branch_name,
            project_dir=project_dir,
            target_branch=target_branch,
        )

        if merge_result.status != MergeStatus.SUCCESS:
            # Merge itself failed (conflict / lock / git error) — fail the task.
            reason = merge_result.message or f"Merge failed: {merge_result.status.value}"
            failed_task = await store.fail_task(pool, task_id, reason)
            await cache.sync_task(redis_client, failed_task)
            await publish_event(
                redis_client, project_id, EventType.MERGE_FAILED,
                task_id=task_id,
                payload=merge_result.to_dict(),
            )
            return merge_result

        # Merge succeeded — mark that cleanup will be needed.
        cleanup_needed = True
        await _set_merge_phase(pool, task_id, MergePhase.MERGED)

        # ----- Phase 2: Test -----
        await _set_merge_phase(pool, task_id, MergePhase.TESTS_RUNNING)
        heartbeat_task = asyncio.create_task(
            _heartbeat_loop(pool, task_id, interval=120),
            name=f"heartbeat-{task_id}",
        )
        try:
            tests_passed, test_output = await _run_tests(project_dir, test_cmd)
        finally:
            heartbeat_task.cancel()
            try:
                await heartbeat_task
            except asyncio.CancelledError:
                pass

        if not tests_passed:
            # Roll back the merge commit.
            log.warning("Tests failed — resetting merge commit")
            await _run_git("reset", "--hard", "HEAD~1", cwd=project_dir, check=False)

            result = MergeResult(
                status=MergeStatus.TEST_FAILED,
                branch_name=branch_name,
                message="Post-merge tests failed",
                test_output=test_output,
            )
            reason = f"Post-merge tests failed:\n{test_output}"
            failed_task = await store.fail_task(pool, task_id, _truncate_output(reason))
            await cache.sync_task(redis_client, failed_task)
            await publish_event(
                redis_client, project_id, EventType.MERGE_FAILED,
                task_id=task_id,
                payload=result.to_dict(),
            )
            return result

        # ----- Phase 3: Success path — push, mark done -----
        await _set_merge_phase(pool, task_id, MergePhase.TESTS_PASSED)
        await _push_to_origin(project_dir)
        await _set_merge_phase(pool, task_id, MergePhase.PUSHED)

        merge_result.test_output = test_output

        completed = await store.complete_task(
            pool, task_id, merge_result.to_dict(),
        )
        # complete_task may return a single Task or a tuple — we only need the
        # first element for cache sync.
        done_task = completed[0] if isinstance(completed, tuple) else completed
        await cache.sync_task(redis_client, done_task)
        await publish_event(
            redis_client, project_id, EventType.MERGE_SUCCEEDED,
            task_id=task_id,
            payload=merge_result.to_dict(),
        )

        log.info(
            "Merge task %s completed: branch=%s sha=%s",
            task_id, branch_name, merge_result.commit_sha,
        )
        return merge_result

    except Exception as exc:
        # Unexpected error — best-effort fail the task.
        log.exception("Unexpected error in execute_merge_task for %s", task_id)
        try:
            reason = f"Unexpected error during merge: {exc}"
            failed_task = await store.fail_task(pool, task_id, _truncate_output(reason))
            await cache.sync_task(redis_client, failed_task)
            await publish_event(
                redis_client, project_id, EventType.MERGE_FAILED,
                task_id=task_id,
                payload={"error": str(exc), "branch_name": branch_name},
            )
        except Exception:
            log.error("Failed to mark task %s as failed", task_id, exc_info=True)
        return MergeResult(
            status=MergeStatus.GIT_ERROR,
            branch_name=branch_name,
            message=f"Unexpected error: {exc}",
        )

    finally:
        # ----- Cleanup: worktree + branch (runs in both success & failure) -----
        if cleanup_needed:
            try:
                await _cleanup_worktree(project_dir, worktree_path, branch_name)
                await _set_merge_phase(pool, task_id, MergePhase.CLEANED_UP)
            except Exception:
                log.warning("Worktree cleanup failed (non-fatal)", exc_info=True)
