"""
Enhanced asynchronous branch manager for handling Looker Git operations.
"""

import asyncio
import hashlib
import logging
import time
from dataclasses import dataclass
from typing import Dict, List, Optional, Set

from looker_validator.async_client import AsyncLookerClient
from looker_validator.exceptions import BranchError, LookerApiError, LookerValidatorException

logger = logging.getLogger(__name__)


@dataclass
class ProjectState:
    """State of a Looker project's Git branch."""
    project: str
    workspace: str
    branch: str
    commit: str


class AsyncBranchManager:
    """Async manager for Git branches in Looker projects.
    
    Handles branch checkout, creation, deletion, and state management.
    """
    
    def __init__(
        self,
        client: AsyncLookerClient,
        project: str,
        remote_reset: bool = False,
        pin_imports: Optional[Dict[str, str]] = None,
        use_personal_branch: bool = False,
        skip_imports: Optional[List[str]] = None,
        use_ci_branch: bool = False,
        ci_run_id: Optional[str] = None,
    ):
        """Initialize the branch manager.

        Args:
            client: AsyncLookerClient instance
            project: Looker project name
            remote_reset: Whether to reset to the remote branch state
            pin_imports: Dict of project:ref pairs for imported projects
            use_personal_branch: Whether to use personal branch
            skip_imports: List of projects to skip importing
            use_ci_branch: Whether to create a unique temporary branch per CI run
            ci_run_id: Unique CI run identifier for branch naming
        """
        self.client = client
        self.project = project
        self.remote_reset = remote_reset
        self.pin_imports = pin_imports or {}
        self.history: List[ProjectState] = []

        self.commit: Optional[str] = None
        self.branch: Optional[str] = None
        self.is_temp_branch: bool = False
        self.use_personal_branch: bool = use_personal_branch
        self.personal_branch: Optional[str] = None
        self.import_managers: List[AsyncBranchManager] = []
        self.skip_imports: List[str] = [] if skip_imports is None else skip_imports
        self.processed_imports: Set[str] = set()
        self.use_ci_branch: bool = use_ci_branch
        self.ci_run_id: Optional[str] = ci_run_id
    
    async def __aenter__(self) -> "AsyncBranchManager":
        """Set up Git branch state."""
        logger.debug(f"Setting up branch manager for {self.project}")
        
        # Store initial state
        state = await self.get_project_state()
        self.history = [state]
        
        # Set up imported projects if in dev mode
        if state.workspace == "dev" and self.pin_imports:
            await self._setup_imported_projects()
        
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Clean up Git branch state."""
        logger.debug(f"Cleaning up branch manager for {self.project}")
        
        # Clean up import projects first
        for manager in self.import_managers:
            await manager.__aexit__(exc_type, exc_val, exc_tb)
        
        # Delete temp branch if created
        if self.is_temp_branch and self.branch:
            await self._cleanup_temp_branch()
        
        # Restore initial state
        await self._restore_initial_state()
    
    async def get_project_state(self) -> ProjectState:
        """Get the current project state.
        
        Returns:
            ProjectState with workspace, branch, and commit
        """
        workspace = await self.client.get_workspace()
        branch_info = await self.client.get_active_branch(self.project)
        
        return ProjectState(
            project=self.project,
            workspace=workspace,
            branch=branch_info["name"],
            commit=branch_info["ref"]
        )
    
    async def _restore_initial_state(self) -> None:
        """Restore the initial workspace and branch."""
        if not self.history:
            return
            
        initial_state = self.history[0]
        
        if initial_state.workspace == "production":
            await self.client.update_workspace("production")
        else:
            await self.client.update_workspace("dev")
            if initial_state.branch:
                await self.client.checkout_branch(self.project, initial_state.branch)
    
    async def setup_branch(
        self, 
        branch: Optional[str] = None, 
        commit_ref: Optional[str] = None
    ) -> None:
        """Set up branch for validation.
    
        Args:
            branch: Git branch name (None for production)
            commit_ref: Git commit reference
        """
        logger.debug(f"Setting up branch for validation: branch={branch}, commit_ref={commit_ref}")
    
        self.branch = None
        self.commit = None
    
        # Branch was specified, use dev mode and checkout the branch
        if branch:
            self.branch = branch
            await self.client.update_workspace("dev")

            if self.use_ci_branch:
                self.branch = await self._checkout_ci_branch(branch)
            elif self.use_personal_branch:
                # Use personal branch, reset to the PR branch
                self.branch = await self._checkout_personal_branch(branch)
            else:
                # Regular checkout
                await self.client.checkout_branch(self.project, branch)

                # Reset to remote if requested
                if self.remote_reset:
                    logger.debug(f"Resetting branch '{branch}' to remote")
                    await self.client.reset_to_remote(self.project)

            # Verify checkout
            await self.verify_branch_checkout()
    
        # Commit reference was specified, create temp branch
        elif commit_ref:
            self.commit = commit_ref
            self.branch = await self._checkout_temp_branch(commit_ref)
    
        # Neither branch nor commit ref was specified, use production
        else:
            await self.client.update_workspace("production")
    
    async def _get_project_imports(self) -> List[str]:
        """Get list of imported projects.
        
        Returns:
            List of project names
        """
        try:
            manifest = await self.client.get_manifest(self.project)
            return [p["name"] for p in manifest.get("imports", []) if not p.get("is_remote", False)]
        except Exception:
            return []  # No imports or error getting manifest
    
    async def _setup_imported_projects(self) -> None:
        """Set up branches for imported projects."""
        # Get all imports
        imports = await self._get_project_imports()
        
        if not imports:
            logger.debug(f"Project '{self.project}' has no imported projects")
            return
            
        logger.debug(f"Project '{self.project}' imports: {imports}")
        
        # Process only imports with pins
        for import_project in imports:
            # Skip if we've already processed this import
            if import_project in self.processed_imports or import_project in self.skip_imports:
                logger.debug(f"Skipping already processed import: {import_project}")
                continue
                
            # If project is pinned, set up branch manager for it
            if import_project in self.pin_imports:
                ref = self.pin_imports[import_project]
                logger.debug(f"Setting up branch for imported project '{import_project}' @ {ref}")
                
                # Create branch manager for import
                import_manager = AsyncBranchManager(
                    self.client,
                    import_project,
                    remote_reset=self.remote_reset,
                    pin_imports=self.pin_imports,
                    use_personal_branch=self.use_personal_branch,
                    skip_imports=self.skip_imports + [self.project],  # Prevent circular imports
                    use_ci_branch=self.use_ci_branch,
                    ci_run_id=self.ci_run_id,
                )
                
                # Enter context and set up branch
                await import_manager.__aenter__()
                is_commit = ref and ref.startswith("refs/") or (len(ref) >= 7 and all(c in "0123456789abcdef" for c in ref))
                if is_commit:
                    await import_manager.setup_branch(commit_ref=ref)
                else:
                    await import_manager.setup_branch(branch=ref)
                
                # Add to managed imports
                self.import_managers.append(import_manager)
                
            # Mark as processed to avoid circular imports
            self.processed_imports.add(import_project)
    
    async def _checkout_ci_branch(self, ref: str) -> str:
        """Create and checkout a unique CI branch from the given ref.

        This avoids race conditions when multiple CI runs share the same
        service account by giving each run its own temporary branch.

        Args:
            ref: Git reference (branch name or SHA) to base the CI branch on

        Returns:
            The CI branch name
        """
        # Step 1: Switch to dev workspace
        logger.info(f"[ci-branch] Step 1: Switching to dev workspace")
        await self.client.update_workspace("dev")

        # Step 2: Resolve ref to SHA
        commit_sha = await self._resolve_ref_to_sha(ref)
        logger.info(f"[ci-branch] Step 2: Resolved '{ref}' to SHA '{commit_sha}'")

        # Step 3: Generate branch name
        branch_name = f"ci-validation-{self.ci_run_id}"
        logger.info(f"[ci-branch] Step 3: Using branch name '{branch_name}'")

        # Step 4: Try deleting existing branch with same name (handles re-runs)
        try:
            await self.client.delete_branch(self.project, branch_name)
            logger.info(f"[ci-branch] Step 4: Deleted existing branch '{branch_name}'")
        except Exception:
            logger.debug(f"[ci-branch] Step 4: No existing branch '{branch_name}' to delete")

        # Step 5: Create branch from SHA
        logger.info(f"[ci-branch] Step 5: Creating branch '{branch_name}' from '{commit_sha}'")
        await self.client.create_branch(self.project, branch_name, commit_sha)

        # Step 6: Checkout branch
        logger.info(f"[ci-branch] Step 6: Checking out branch '{branch_name}'")
        await self.client.checkout_branch(self.project, branch_name)

        self.is_temp_branch = True
        logger.info(f"[ci-branch] Successfully set up CI branch '{branch_name}' from '{ref}' ({commit_sha})")
        return branch_name

    async def _checkout_personal_branch(self, ref: str) -> str:
        """Updates the user's personal branch to the git ref.

        Args:
            ref: Git reference to reset to

        Returns:
            Branch name
        """
        # Step 1: Switch to dev workspace
        logger.info(f"[personal-branch] Step 1: Switching to dev workspace")
        await self.client.update_workspace("dev")

        # Step 2: Find personal branch
        if not self.personal_branch:
            self.personal_branch = await self._get_personal_branch()
        logger.info(f"[personal-branch] Step 2: Found personal branch '{self.personal_branch}'")

        # Step 3: Checkout personal branch
        logger.info(f"[personal-branch] Step 3: Checking out personal branch '{self.personal_branch}'")
        await self.client.checkout_branch(self.project, self.personal_branch)

        # Step 4: Reset to remote (may fail if no upstream, that's OK)
        try:
            logger.info(f"[personal-branch] Step 4: Resetting to remote")
            await self.client.reset_to_remote(self.project)
        except Exception as e:
            logger.warning(f"[personal-branch] Step 4: reset_to_remote failed (non-fatal): {e}")

        # Step 5: Resolve the target branch name to a commit SHA
        commit_sha = await self._resolve_ref_to_sha(ref)
        logger.info(f"[personal-branch] Step 5: Resolved '{ref}' to ref '{commit_sha}'")

        # Step 6: Hard reset to the resolved commit SHA
        logger.info(f"[personal-branch] Step 6: Hard resetting '{self.personal_branch}' to '{commit_sha}'")
        await self.client.hard_reset_branch(self.project, self.personal_branch, commit_sha)

        logger.info(f"[personal-branch] Successfully set up personal branch '{self.personal_branch}' -> '{ref}' ({commit_sha})")
        return self.personal_branch
    
    async def verify_branch_checkout(self) -> bool:
        """Verify that the current branch is what we expect.
    
        Returns:
            True if correct branch is checked out
        """
        if not self.branch:
            return True
        
        try:
            current_branch = await self.client.get_active_branch_name(self.project)
            if current_branch != self.branch:
                logger.warning(f"Branch mismatch: Expected '{self.branch}', found '{current_branch}'")
            
                # Check out the correct branch again
                await self.client.checkout_branch(self.project, self.branch)
            
                # Verify again
                current_branch = await self.client.get_active_branch_name(self.project)
                if current_branch != self.branch:
                    logger.error(f"Failed to checkout branch '{self.branch}', still on '{current_branch}'")
                    return False
            
                logger.info(f"Successfully checked out branch '{self.branch}'")
            
            return True
        except Exception as e:
            logger.error(f"Error verifying branch: {str(e)}")
            return False
    
    async def _resolve_ref_to_sha(self, ref: str) -> str:
        """Resolve a branch name or ref to a commit SHA.

        The Looker API's hard_reset_branch expects a commit SHA, not a branch name.
        This method looks up the branch and returns its commit SHA.

        Args:
            ref: Branch name or git reference

        Returns:
            Commit SHA string
        """
        # If it already looks like a SHA (hex string, 7+ chars), return as-is
        if len(ref) >= 7 and all(c in "0123456789abcdef" for c in ref):
            return ref

        # Look up the branch to get its commit SHA.
        # Looker's branch list includes both local and remote branches.
        # Remote branches that were never locally checked out may only appear
        # with an "origin/" prefix, so check both formats.
        branches = await self.client.get_all_branches(self.project)
        for branch in branches:
            branch_name = branch.get("name", "")
            if branch_name == ref or branch_name == f"origin/{ref}":
                sha = branch.get("ref", "")
                if sha:
                    logger.debug(f"Resolved branch '{ref}' (listed as '{branch_name}') to SHA '{sha}'")
                    return sha

        # Branch not found in Looker's local list — it's likely a remote branch
        # that was never checked out in this Looker workspace.
        # Use the remote tracking ref (origin/<branch>) directly so that git
        # can resolve it to the correct SHA without a local checkout.
        remote_ref = f"origin/{ref}"
        logger.warning(
            f"Branch '{ref}' not found in Looker's local branch list. "
            f"Falling back to remote tracking ref '{remote_ref}'."
        )
        return remote_ref

    async def _get_personal_branch(self) -> str:
        """Finds the name of the user's personal branch.
        
        Returns:
            Personal branch name
        """
        branches = await self.client.get_all_branches(self.project)
        for branch in branches:
            if branch.get("personal") and not branch.get("readonly"):
                return branch["name"]
                
        raise BranchError(
            title=f"Personal branch not found",
            detail=f"Could not find a personal branch for project '{self.project}'"
        )
    
    async def _checkout_temp_branch(self, ref: str) -> str:
        """Creates a temporary branch off a commit or production.
        
        Args:
            ref: Git reference to base branch on
            
        Returns:
            Temporary branch name
        """
        # Save the dev mode state
        await self.client.update_workspace("dev")
        self.history.append(await self.get_project_state())
        
        # Generate a unique branch name
        branch_name = f"tmp_validator_{self._generate_hash()}"
        
        logger.debug(f"Creating temporary branch '{branch_name}' from ref '{ref}'")
        
        # Create and checkout the branch
        await self.client.create_branch(self.project, branch_name, ref)
        await self.client.checkout_branch(self.project, branch_name)
        
        self.is_temp_branch = True
        return branch_name
    
    async def _cleanup_temp_branch(self) -> None:
        """Clean up the temporary branch."""
        if not self.is_temp_branch or not self.branch:
            return
            
        try:
            logger.debug(f"Cleaning up temporary branch '{self.branch}'")
            
            # We need to be in dev mode to delete branches
            await self.client.update_workspace("dev")
            
            # Switch to the initial branch before deleting
            dev_state = self.history[1] if len(self.history) > 1 else None
            if dev_state and dev_state.branch:
                await self.client.checkout_branch(self.project, dev_state.branch)
                
            # Delete the temporary branch
            await self.client.delete_branch(self.project, self.branch)
            
            self.is_temp_branch = False
        except Exception as e:
            logger.warning(f"Failed to clean up temporary branch: {str(e)}")
    
    def _generate_hash(self) -> str:
        """Generate a short hash for temporary branch names."""
        hash_obj = hashlib.sha1()
        hash_obj.update(str(time.time()).encode('utf-8'))
        return hash_obj.hexdigest()[:8]