from typing import List, Optional

import click
from loguru import logger
from pydantic import BaseModel

import ghchain
from ghchain.commit import Commit
from ghchain.git_utils import (
    create_branch_name,
    get_all_branches,
    get_current_branch,
    git_push,
)
from ghchain.github_utils import (
    create_branch_from_issue,
    create_pull_request,
    get_next_gh_id,
    run_tests_on_branch,
    update_pr_descriptions,
)
from ghchain.utils import run_command


class Stack(BaseModel):
    commits: List[Commit]
    dev_branch: str
    base_branch: str

    @classmethod
    def create(cls, base_branch: Optional[str] = None) -> "Stack":
        """
        Create a Stack object with commits from the current branch which are not in the base branch.
        Whatever branch is currently checked out will be considered the dev branch.
        """
        if not base_branch:
            base_branch = ghchain.config.base_branch
        dev_branch = get_current_branch()

        logger.debug(f"Creating stack with dev branch: {dev_branch}")

        # Fetch the commits that are in dev_branch but not in base_branch
        commits_diff = ghchain.repo.iter_commits(f"{base_branch}..{dev_branch}")

        commits = []
        # Reverse log output to start from the bottom of the stack
        for commit in reversed(list(commits_diff)):
            sha = commit.hexsha
            message = commit.message.strip()
            is_fixup = message.startswith("fixup!")

            # Find branches that point to this commit (branches where this commit is the HEAD)
            pointing_branches = [
                ref.name.split("/")[-1]
                for ref in ghchain.repo.branches
                if ref.commit == commit
            ]
            pointing_branches = set(pointing_branches) - {
                dev_branch
            }  # Exclude the dev branch

            if len(pointing_branches) > 1:
                error_message = f"Commit {sha} has multiple branches: {pointing_branches}. This is not supported."
                logger.error(error_message)
                raise ValueError(error_message)

            branch = pointing_branches.pop() if pointing_branches else None
            commit_obj = Commit(
                sha=sha,
                branch=branch,
                message=message,
                is_fixup=is_fixup,
            )
            commits.append(commit_obj)

            if is_fixup:
                # Find the previous commit that is not a fixup and assign it the same branch
                target_commit = next(
                    (
                        commit
                        for commit in commits[::-1]
                        if not commit.is_fixup and not commit.branch
                    ),
                    None,
                )
                if target_commit:
                    target_commit.branch = commit_obj.branch

        return cls(commits=commits, dev_branch=dev_branch, base_branch=base_branch)

    @property
    def commit2idx(self):
        return {commit.sha: i for i, commit in enumerate(self.commits)}

    @property
    def branch_ids(self) -> set[int]:
        return {
            int(commit.branch.split("-")[-1])
            for commit in self.commits
            if commit and commit.branch
            if commit.branch.split("-")[-1].isnumeric()
        }

    @property
    def branches(self) -> list[str]:
        """
        Return the branches in the stack, sorted by order in the stack.
        """
        return [commit.branch for commit in self.commits if commit.branch]

    def process_commit(
        self,
        commit: Commit,
        create_pr: bool = False,
        draft: bool = False,
        with_tests: bool = False,
    ):
        """
        Create a branch for each commit in the stack that doesn't already have one.
        If create_pr is True, create a PR for each branch.
        """
        pr_created = False
        if not commit.branch:
            if commit.issue_url:
                logger.info(f"Creating branch from issue {commit.issue_url}")
                # set a temporary branch name as it will be created later
                branch_name = "{will be created from issue}"

            else:
                branch_id = max([get_next_gh_id(), *[id + 1 for id in self.branch_ids]])
                branch_name = create_branch_name(
                    ghchain.config.branch_name_template, branch_id
                )
            click.confirm(
                f"Create branch {branch_name} for commit {commit.sha}?\n{commit.message}\n",
                abort=True,
                default=True,
            )
            logger.info(f"Creating branch {branch_name} for commit {commit.sha}")
            if commit.issue_url:
                # create the branch from the gh issue
                branch_name = create_branch_from_issue(
                    issue_id=commit.issue_url, base_commit=commit.sha
                )
            else:
                ghchain.repo.git.branch(branch_name, commit.sha)
            commit.branch = branch_name

            # Push the branch to the remote
            git_push(branch_name)

        else:
            branch_name = commit.branch

        # If the commit is not a fixup and it does not have a PR, create one
        if not commit.is_fixup and create_pr and not commit.pr_id:
            commit.pr_url = create_pull_request(
                base_branch=ghchain.config.base_branch.replace("origin/", "")
                if not self.commit2idx[commit.sha]
                else self.commits[self.commit2idx[commit.sha] - 1].branch,
                head_branch=branch_name,
                title=commit.message.split("\n")[0],
                body=commit.message,
                draft=draft,
            )
            pr_created = True

        if create_pr:
            update_pr_descriptions(
                pr_stack=[
                    c.pr_url
                    for c in self.commits[: self.commit2idx[commit.sha] + 1]
                    if c.pr_url
                ]
            )

        if with_tests:
            run_tests_on_branch(branch_name, commit.pr_url)

        return pr_created

    def publish(self):
        """
        Collect all branches in the stack that need to be pushed to the remote,
        ask the user for confirmation, and then push them all with force-with-lease.
        """
        branches_to_push = []

        # Collect branches that need to be pushed
        for commit in self.commits:
            if not commit.branch:
                logger.trace(
                    f"Commit {commit.sha} does not have an associated branch. Skipping."
                )
                continue

            branch_name = commit.branch
            remote_branch = f"origin/{branch_name}"

            try:
                local_sha = ghchain.repo.git.rev_parse(branch_name)
                remote_sha = ghchain.repo.git.rev_parse(remote_branch)

                if local_sha != remote_sha:
                    branches_to_push.append(branch_name)
            except Exception:
                # Remote branch does not exist
                branches_to_push.append(branch_name)

        if not branches_to_push:
            logger.info("All branches are up-to-date with the remote. Nothing to push.")
            return

        # Confirm with the user
        click.confirm(
            f"The following branches will be pushed with --force-with-lease:\n{', '.join(branches_to_push)}\n"
            "Do you want to proceed?",
            abort=True,
        )

        # Push all branches with --force-with-lease
        try:
            ghchain.repo.git.push("--force-with-lease", "origin", *branches_to_push)
            logger.info(
                f"Successfully pushed branches: {', '.join(branches_to_push)} with --force-with-lease."
            )
        except Exception as push_error:
            logger.error(
                f"Failed to push branches: {', '.join(branches_to_push)}: {push_error}"
            )

        logger.info("Publishing complete.")


def find_branch_of_rebased_commit(commit: Commit) -> str:
    """
    When a commit is rebased, the commit hash changes. This function checks if there is any branch
    that has a commit with the message as top commit. If there is, it returns the branch name.
    """
    # Get the commit message
    commit_message = commit.message

    # List all branches
    result = get_all_branches()
    branches = result.stdout.splitlines()

    # Check the top commit message of each branch
    for branch in branches:
        branch = branch.strip().replace(
            "* ", ""
        )  # Remove leading '*' for the current branch
        result = run_command(
            ["git", "log", "-1", "--pretty=%B", branch],
            check=True,
        )
        top_commit_message = result.stdout.strip()
        if top_commit_message == commit_message:
            logger.debug(f"Found branch {branch} with commit message {commit_message}")
            return branch

    return None
