import asyncio
import os
import pathlib
import random
import shlex
import subprocess
import time
import uuid
from argparse import Namespace
from typing import Any, cast

import httpx
import requests

import firebase_admin
from firebase_admin import credentials, firestore

from resolver.daytona_patch import apply_daytona_patch
from resolver.error_tracker import ErrorTracker
from resolver.resolver_output import CustomResolverOutput
from resolver.secrets import Secrets
from resolver.send_pull_request import (
    apply_patch,
    initialize_repo,
    make_commit_with_summary,
)
from resolver.utils import get_comprehensive_language_info, load_firebase_config

# Apply daytona compatibility patch before any openhands imports
apply_daytona_patch()

from openhands.core.config import LLMConfig  # noqa: E402
from openhands.core.logger import openhands_logger as logger  # noqa: E402
from openhands.integrations.service_types import ProviderType  # noqa: E402
from openhands.resolver.interfaces.issue import Issue  # noqa: E402
from openhands.resolver.issue_handler_factory import IssueHandlerFactory  # noqa: E402
from openhands.resolver.issue_resolver import IssueResolver  # noqa: E402
from openhands.resolver.resolver_output import ResolverOutput  # noqa: E402
from openhands.runtime import Runtime  # noqa: E402

# Monkey patch OpenHands LLM to handle GPT-5 models
def patch_openhands_for_gpt5():
    """Patch OpenHands to exclude GPT-5 models from receiving stop tokens."""
    try:
        from openhands.llm import llm as openhands_llm_module
        if hasattr(openhands_llm_module, 'MODELS_WITHOUT_STOP_WORDS'):
            # Add GPT-5 models to the exclusion list (both base and full proxy names)
            gpt5_models = ['gpt-5', 'litellm_proxy/neulab/gpt-5']
            current_models = list(openhands_llm_module.MODELS_WITHOUT_STOP_WORDS)
            for gpt5_model in gpt5_models:
                if gpt5_model not in current_models:
                    current_models.append(gpt5_model)
            openhands_llm_module.MODELS_WITHOUT_STOP_WORDS = current_models
            # logger.info(f"Updated MODELS_WITHOUT_STOP_WORDS to include GPT-5: {openhands_llm_module.MODELS_WITHOUT_STOP_WORDS}")
    except Exception as e:
        logger.warning(f"Failed to patch OpenHands for GPT-5: {e}")

# Apply the patch
patch_openhands_for_gpt5()

# Don't make this configurable for now, unless we have other competitive agents
AGENT_CLASS = "CodeActAgent"


class PRArenaIssueResolver(IssueResolver):
    """PR Arena Issue Resolver that processes issues with multiple LLM models."""

    # Type annotations for class attributes
    llm_config: LLMConfig
    issue_handler: Any  # This will be set dynamically in _resolve_with_model
    output_dir: str

    def __init__(self, args: Namespace) -> None:
        # super().__init__(args) # Most shared arguments are processed by parent class

        # Setup and validate container images
        self.sandbox_config = self._setup_sandbox_config(
            args.base_container_image,
            args.runtime_container_image,
            args.is_experimental,
        )

        parts = args.selected_repo.rsplit("/", 1)
        if len(parts) < 2:
            raise ValueError("Invalid repository format. Expected owner/repo")
        owner, repo = parts

        token = args.token or os.getenv("GITHUB_TOKEN") or os.getenv("GITLAB_TOKEN")
        username = args.username if args.username else os.getenv("GIT_USERNAME")
        if not username:
            raise ValueError("Username is required.")

        if not token:
            raise ValueError("Token is required.")

        platform = ProviderType.GITHUB

        base_url = args.llm_base_url
        api_version = os.environ.get("LLM_API_VERSION", None)
        llm_num_retries = int(os.environ.get("LLM_NUM_RETRIES", "4"))
        llm_retry_min_wait = int(os.environ.get("LLM_RETRY_MIN_WAIT", "5"))
        llm_retry_max_wait = int(os.environ.get("LLM_RETRY_MAX_WAIT", "30"))
        llm_retry_multiplier = int(os.environ.get("LLM_RETRY_MULTIPLIER", 2))
        llm_timeout = int(os.environ.get("LLM_TIMEOUT", 0))

        # Initialize values for custom resolver
        self.token = token
        self.username = username
        Secrets.TOKEN = self.token

        multiple_models = args.llm_models or os.environ["LLM_MODELS"]
        if multiple_models:
            model_names = [model.strip() for model in multiple_models.split(",")]
        else:
            raise ValueError(
                "No LLM models provided in either the arguments or environment variables."
            )

        selected_models = random.sample(model_names, 2)
        # selected_models = model_names
        
        # Save selected models to file for timeout handling
        with open('/tmp/selected_models.txt', 'w') as f:
            f.write(','.join(selected_models))
        
        self.llm_configs = []

        for model in selected_models:
            # Determine if this model needs special parameter handling
            needs_drop_params = (
                "o1-mini" in model
                or "o3-mini" in model
                or "o4-mini" in model
                or "gemini" in model.lower()
                or "gpt-5" in model.lower()
            )

            if "gpt-5" in model.lower():
                gpt5_timeout = 300  # 5 minutes per API call
                print(f"GPT-5 detected: {model}, applying specialized configuration")

                # GPT-5 needs very specific configuration - only default temperature (1.0) is supported
                config_params = {
                    "model": model,
                    "api_key": Secrets.get_api_key(),
                    "base_url": base_url,
                    "num_retries": llm_num_retries,
                    "retry_min_wait": llm_retry_min_wait,
                    "retry_max_wait": llm_retry_max_wait,
                    "retry_multiplier": llm_retry_multiplier,
                    "timeout": gpt5_timeout,
                    "drop_params": True,  # Critical for GPT-5
                    # Note: GPT-5 only supports default temperature (1.0)
                }
                
                llm_config = LLMConfig(**config_params)
                
                # GPT-5 specific parameter cleanup - remove all custom parameters that aren't supported
                if hasattr(llm_config, "temperature"):
                    llm_config.temperature = None  # Use default (1.0)
                if hasattr(llm_config, "stop"):
                    llm_config.stop = None
                if hasattr(llm_config, "top_p"):
                    llm_config.top_p = None
                if hasattr(llm_config, "max_tokens"):
                    llm_config.max_tokens = 2000  # Shorter responses for precision
                if hasattr(llm_config, "frequency_penalty"):
                    llm_config.frequency_penalty = None
                if hasattr(llm_config, "presence_penalty"):
                    llm_config.presence_penalty = None
            else:
                llm_config = LLMConfig(
                    model=model,
                    api_key=Secrets.get_api_key(),
                    base_url=base_url,
                    num_retries=llm_num_retries,
                    retry_min_wait=llm_retry_min_wait,
                    retry_max_wait=llm_retry_max_wait,
                    retry_multiplier=llm_retry_multiplier,
                    timeout=llm_timeout,
                    drop_params=needs_drop_params,
                )

            if "gemini" in model.lower():
                # Gemini models have specific limitations on tool formats
                # Set additional Gemini-specific configurations if needed
                if hasattr(llm_config, "supports_function_calling"):
                    llm_config.supports_function_calling = True
                # Gemini models work better with simplified tool schemas
                if hasattr(llm_config, "simplify_tools"):
                    llm_config.simplify_tools = True

            if "o4-mini" in model and hasattr(llm_config, "top_p"):
                # o4-mini models require top_p to be set to None
                llm_config.top_p = None

            if "o3-mini" in model and hasattr(llm_config, "top_p"):
                # o3-mini models require top_p to be set to None
                llm_config.top_p = None

            if "o1-mini" in model and hasattr(llm_config, "top_p"):
                # o1-mini models require top_p to be set to None
                llm_config.top_p = None

            if "o1-mini" in model and hasattr(llm_config, "stop"):
                # o1-mini models require stop to be set to None
                llm_config.stop = None

            self.llm_configs.append(llm_config)

            # Only set api_version if it was explicitly provided, otherwise let LLMConfig handle it
            if api_version is not None:
                llm_config.api_version = api_version

        repo_instruction = None
        if args.repo_instruction_file:
            with open(args.repo_instruction_file, "r") as f:
                repo_instruction = f.read()

        issue_type = args.issue_type

        # Read the prompt template
        prompt_file = os.path.join(
            os.path.dirname(__file__), "prompts/resolve/basic-with-tests.jinja"
        )

        with open(prompt_file, "r") as f:
            user_instructions_prompt_template = f.read()

        with open(
            prompt_file.replace(".jinja", "-conversation-instructions.jinja")
        ) as f:
            conversation_instructions_prompt_template = f.read()

        self.owner = owner
        self.repo = repo
        self.platform = platform
        self.max_iterations = args.max_iterations
        self.user_instructions_prompt_template = user_instructions_prompt_template
        self.conversation_instructions_prompt_template = (
            conversation_instructions_prompt_template
        )
        self.issue_type = issue_type
        self.repo_instruction = repo_instruction
        self.issue_number = args.issue_number
        self.comment_id = args.comment_id

        raw_config = Secrets.get_firebase_config()
        self.firebase_config = load_firebase_config(raw_config)

        # Initialize error tracker
        self.error_tracker = ErrorTracker(
            owner=self.owner,
            repo=self.repo,
            issue_number=self.issue_number,
            token=self.token,
        )

    async def complete_runtime(
        self,
        runtime: Runtime,
        base_commit: str,
    ) -> dict[str, Any]:
        patch = await super().complete_runtime(runtime, base_commit)
        runtime.close()
        return patch

    async def resolve_issues_with_random_models(self) -> None:
        uuid_ref = None
        models_info = {
            "model1": self.llm_configs[0].model.split("/")[-1]
            if len(self.llm_configs) > 0
            else "unknown",
            "model2": self.llm_configs[1].model.split("/")[-1]
            if len(self.llm_configs) > 1
            else "unknown",
        }

        try:
            # Run both agents in parallel using asyncio.gather
            # return_exceptions=True ensures we can handle individual failures
            results = await asyncio.gather(
                self._resolve_with_model(model_index=0, output_dir="output1"),
                self._resolve_with_model(model_index=1, output_dir="output2"),
                return_exceptions=True,
            )

            # Handle exceptions from either task
            resolver_output_1 = results[0]
            resolver_output_2 = results[1]

            if isinstance(resolver_output_1, Exception):
                raise resolver_output_1
            if isinstance(resolver_output_2, Exception):
                raise resolver_output_2

            # Send both outputs to Firebase after both complete
            # Type assertions to help mypy understand these are CustomResolverOutput
            assert isinstance(resolver_output_1, CustomResolverOutput)
            assert isinstance(resolver_output_2, CustomResolverOutput)

            await self.send_to_firebase(
                resolved_output_1=resolver_output_1,
                resolved_output_2=resolver_output_2,
                pr_type="draft",
            )

            # Get UUID from environment for error tracking if needed
            uuid_ref = os.getenv("UUID")

        except Exception as e:
            logger.error(f"Error in resolve_issues_with_random_models: {str(e)}")

            # Log the error to error_collection
            try:
                # Check if this was a git patch empty error
                git_patch_empty = "Openhands failed to make code changes" in str(e)

                # Determine which model caused the error
                model_error_details = None
                if models_info:
                    model_names = list(models_info.values())
                    if model_names:
                        model_error_details = (
                            f"Error occurred with models: {', '.join(model_names)}"
                        )

                await self.error_tracker.log_error(
                    error_type="agent_failure",
                    error_message=f"Agent failed during issue resolution: {str(e)}",
                    uuid_ref=uuid_ref,
                    models=models_info,
                    additional_context={
                        "stage": "issue_resolution",
                        "exception_type": type(e).__name__,
                    },
                    git_patch_empty=git_patch_empty,
                    model_error_details=model_error_details,
                )
            except Exception as tracking_error:
                logger.error(f"Failed to log error to Firebase: {tracking_error}")

            # Set FAILED=TRUE in GitHub environment
            github_env_path = os.getenv("GITHUB_ENV")
            if github_env_path:
                try:
                    with open(github_env_path, "a") as env_file:
                        env_file.write("FAILED=TRUE\n")
                except Exception as env_error:
                    logger.error(f"Failed to write to GITHUB_ENV: {env_error}")

            # Re-raise the exception so workflow can handle it
            raise

    async def _check_and_log_empty_patches(
        self,
        resolved_output_1: CustomResolverOutput,
        resolved_output_2: CustomResolverOutput,
        uuid_ref: str,
    ) -> None:
        """Check for empty git patches in successful resolutions and log them."""
        models_info = {
            "model1": resolved_output_1.model
            if hasattr(resolved_output_1, "model") and resolved_output_1.model
            else "unknown",
            "model2": resolved_output_2.model
            if hasattr(resolved_output_2, "model") and resolved_output_2.model
            else "unknown",
        }

        # Check if either model produced an empty patch
        model1_empty = (
            not resolved_output_1.git_patch or resolved_output_1.git_patch.strip() == ""
        )
        model2_empty = (
            not resolved_output_2.git_patch or resolved_output_2.git_patch.strip() == ""
        )

        if model1_empty or model2_empty:
            empty_models = []
            if model1_empty:
                empty_models.append(f"model1 ({models_info['model1']})")
            if model2_empty:
                empty_models.append(f"model2 ({models_info['model2']})")

            error_message = (
                f"One or more models produced empty patches: {', '.join(empty_models)}"
            )
            model_error_details = f"Empty patches from: {', '.join(empty_models)}"

            try:
                await self.error_tracker.log_error(
                    error_type="empty_patch_success",
                    error_message=error_message,
                    uuid_ref=uuid_ref,
                    models=models_info,
                    additional_context={
                        "stage": "successful_resolution_with_empty_patch",
                        "model1_empty": model1_empty,
                        "model2_empty": model2_empty,
                        "both_empty": model1_empty and model2_empty,
                    },
                    git_patch_empty=True,
                    model_error_details=model_error_details,
                )
                logger.info(
                    f"Logged empty patch tracking for successful resolution: {error_message}"
                )
            except Exception as tracking_error:
                logger.error(
                    f"Failed to log empty patch tracking to Firebase: {tracking_error}"
                )

    async def send_to_firebase(
        self,
        resolved_output_1: CustomResolverOutput,
        resolved_output_2: CustomResolverOutput,
        pr_type: str,
    ) -> None:
        """
        Send the resolver output to Firebase Firestore.

        Args:
            resolved_output (ResolverOutput): The resolved output to be sent.
            owner (str): GitHub owner.
            repo (str): GitHub repository name.
            issue_number (int): Issue number.
            firebase_config (dict): Firebase configuration.
        """
        pathlib.Path("output1").mkdir(parents=True, exist_ok=True)
        pathlib.Path("output2").mkdir(parents=True, exist_ok=True)

        file_name = "output.jsonl"
        output_file1 = pathlib.Path("output1") / file_name
        output_file2 = pathlib.Path("output2") / file_name

        # [PR-Arena] Retrieve commit hash and send it to firesbase as well.
        # And somehow save the file somewhere so that send_pull_request.py could get the file (new commit).
        await self.get_new_commit_hash(
            output_dir="output1", resolver_output=resolved_output_1, pr_type=pr_type
        )
        await self.get_new_commit_hash(
            output_dir="output2", resolver_output=resolved_output_2, pr_type=pr_type
        )

        # Write the resolved output to a JSONL file
        with open(output_file1, "a") as output_fp:
            output_fp.write(resolved_output_1.model_dump_json() + "\n")

        with open(output_file2, "a") as output_fp:
            output_fp.write(resolved_output_2.model_dump_json() + "\n")

        # Send the resolved output to Firebase Firestore
        cred = credentials.Certificate(self.firebase_config)
        if not firebase_admin._apps:
            firebase_admin.initialize_app(cred)

        db = firestore.client()

        current_time = firestore.SERVER_TIMESTAMP

        repo_url = f"https://github.com/{self.owner}/{self.repo}"

        # Extract issue details from the first resolver output (both should have the same issue)
        issue = resolved_output_1.issue
        issue_number = issue.number
        issue_title = issue.title
        issue_body = issue.body

        # Collect language information
        language_info = get_comprehensive_language_info(
            owner=self.owner,
            repo=self.repo,
            token=self.token,
            git_patch_1=resolved_output_1.git_patch,
            git_patch_2=resolved_output_2.git_patch,
        )

        # logger.info(f"Language information collected: {language_info}")

        model_reference = {
            "claude-sonnet-4-20250514": "model1",
            "gpt-4.1-2025-04-14": "model2",
            "Meta-Llama-3.1-405B-Instruct": "model3",
            "o4-mini": "model4",
            "gemini-2.5-pro-preview-05-06": "model5",
            "Qwen-QwQ-32B-Preview": "model6",
            "deepseek-v3": "model7",
            "deepseek-r1": "model8",
            "o3-mini": "model9",
            "o1-mini": "model10",
            "kimi-k2-0711-preview": "model11",
            "gemini-2.5-pro-preview-06-05": "model12",
        }

        model1_id = model_reference.get(
            (cast(str, resolved_output_1.model)), "Model ID Not Found"
        )
        model2_id = model_reference.get(
            (cast(str, resolved_output_2.model)), "Model ID Not Found"
        )

        # Create unified issue data structure that handles both empty and non-empty patches
        issue_data = {
            "repo_url": repo_url,
            "issue_number": issue_number,
            "issue_title": issue_title,
            "issue_body": issue_body,
            "owner": self.owner,
            "repo": self.repo,
            "status": "pending",  # Always pending to allow arena comparison
            "language_info": language_info,
            "models": {
                "modelA": {
                    "modelId": model1_id,
                    "modelName": resolved_output_1.model,
                    "commit_hash": resolved_output_1.commit_hash,
                    "agent_code": resolved_output_1.git_patch
                    if resolved_output_1.git_patch
                    else "",
                    "duration": resolved_output_1.duration
                    if resolved_output_1.duration
                    else None,
                    "cost": {
                        "accumulated_cost": resolved_output_1.accumulated_cost,
                        "token_usage": resolved_output_1.token_usage,
                        "cost_per_input_token": resolved_output_1.cost_per_input_token,
                        "cost_per_output_token": resolved_output_1.cost_per_output_token,
                    },
                },
                "modelB": {
                    "modelId": model2_id,
                    "modelName": resolved_output_2.model,
                    "commit_hash": resolved_output_2.commit_hash,
                    "agent_code": resolved_output_2.git_patch
                    if resolved_output_2.git_patch
                    else "",
                    "duration": resolved_output_2.duration
                    if resolved_output_2.duration
                    else None,
                    "cost": {
                        "accumulated_cost": resolved_output_2.accumulated_cost,
                        "token_usage": resolved_output_2.token_usage,
                        "cost_per_input_token": resolved_output_2.cost_per_input_token,
                        "cost_per_output_token": resolved_output_2.cost_per_output_token,
                    },
                },
            },
            "winner": None,  # No winner determined yet
            "createdAt": current_time,
            "updatedAt": current_time,
            "installationToken": self.token,
        }

        reference_id = str(uuid.uuid4())

        issue_ref = db.collection("issue_collection").document(reference_id)
        issue_ref.set(issue_data)

        current_time = firestore.SERVER_TIMESTAMP

        user_data = {
            "githubId": self.owner,
            "createdAt": current_time,
            "lastActive": current_time,
            "selections": {
                reference_id: {
                    "issueId": reference_id,
                    "choice": None,  # No choice made yet
                    "selectedAt": None,
                    "isLatest": True,
                    "language": "en",  # Default language
                    "isAnonymous": True,
                    "deduplicated": True,
                    "programming_language": language_info.get(
                        "primary_language", "Unknown"
                    ),
                    "repo_languages": language_info.get(
                        "repo_language_percentages", {}
                    ),
                    "patch_languages": language_info.get(
                        "patch_language_percentages", {}
                    ),
                    "modelA": {
                        "modelId": model1_id,
                        "modelName": resolved_output_1.model,
                    },
                    "modelB": {
                        "modelId": model2_id,
                        "modelName": resolved_output_2.model,
                    },
                }
            },
        }

        # Store in user_collection with owner as document ID
        user_ref = db.collection("userdata_collection").document(self.owner)
        user_ref.set(user_data, merge=True)

        # Check for empty patches and log them
        await self._check_and_log_empty_patches(
            resolved_output_1, resolved_output_2, reference_id
        )

        github_env_path = os.getenv("GITHUB_ENV")
        if not github_env_path:
            raise RuntimeError("GITHUB_ENV environment variable is not set.")

        # Write the decision to the environment file
        with open(github_env_path, "a") as env_file:
            env_file.write(f"UUID={reference_id}\n")
            env_file.write("FAILED=FALSE\n")

        # print("Data successfully written to Firestore collections 'issue_collection' and 'user_collection'")
        # print(f"Issue ID: {self.issue_number}, Models: {resolved_output_1.model} vs {resolved_output_2.model}")

    async def resolve_issue(
        self,
        reset_logger: bool = False,
    ) -> CustomResolverOutput:
        """Resolve a single issue.

        Args:
            reset_logger: Whether to reset the logger for multiprocessing.
        """
        start_time = time.time()
        output = None
        customOutput = None

        # Load dataset
        issues: list[Issue] = self.issue_handler.get_converted_issues(
            issue_numbers=[self.issue_number], comment_id=self.comment_id
        )

        if not issues:
            raise ValueError(
                f"No issues found for issue number {self.issue_number}. Please verify that:\n"
                f"1. The issue/PR #{self.issue_number} exists in the repository {self.owner}/{self.repo}\n"
                f"2. You have the correct permissions to access it\n"
                f"3. The repository name is spelled correctly"
            )

        issue = issues[0]

        # Update error tracker with issue information
        self.error_tracker.set_issue_info(issue.title, issue.body)

        if self.comment_id is not None:
            if (
                self.issue_type == "pr"
                and not issue.review_comments
                and not issue.review_threads
                and not issue.thread_comments
            ):
                raise ValueError(
                    f"Comment ID {self.comment_id} did not have a match for issue {issue.number}"
                )

            if self.issue_type == "issue" and not issue.thread_comments:
                raise ValueError(
                    f"Comment ID {self.comment_id} did not have a match for issue {issue.number}"
                )

        # Setup directories
        pathlib.Path(self.output_dir).mkdir(parents=True, exist_ok=True)
        pathlib.Path(os.path.join(self.output_dir, "infer_logs")).mkdir(
            parents=True, exist_ok=True
        )
        logger.info(f"Using output directory: {self.output_dir}")

        # checkout the repo
        repo_dir = os.path.join(self.output_dir, "repo")
        if not os.path.exists(repo_dir):
            checkout_output = subprocess.check_output(
                [
                    "git",
                    "clone",
                    self.issue_handler.get_clone_url(),
                    f"{self.output_dir}/repo",
                ]
            ).decode("utf-8")
            if "fatal" in checkout_output:
                raise RuntimeError(f"Failed to clone repository: {checkout_output}")

        # get the commit id of current repo for reproducibility
        base_commit = (
            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=repo_dir)
            .decode("utf-8")
            .strip()
        )
        logger.info(f"Base commit: {base_commit}")

        if self.repo_instruction is None:
            # Check for .openhands_instructions file in the workspace directory
            openhands_instructions_path = os.path.join(
                repo_dir, ".openhands_instructions"
            )
            if os.path.exists(openhands_instructions_path):
                with open(openhands_instructions_path, "r") as f:
                    self.repo_instruction = f.read()

        # OUTPUT FILE
        output_file = os.path.join(self.output_dir, "output.jsonl")
        logger.info(f"Writing output to {output_file}")

        # Check if this issue was already processed
        if os.path.exists(output_file):
            with open(output_file, "r") as f:
                for line in f:
                    data = ResolverOutput.model_validate_json(line)
                    if data.issue.number == self.issue_number:
                        logger.warning(
                            f"Issue {self.issue_number} was already processed. Skipping."
                        )
                        return CustomResolverOutput(**data.model_dump())

        logger.info(
            f"Resolving issue {self.issue_number} with Agent {AGENT_CLASS}, model **MODEL NAME REDACTED**, max iterations {self.max_iterations}."
        )

        try:
            # checkout to pr branch if needed
            if self.issue_type == "pr":
                branch_to_use = issue.head_branch
                logger.info(
                    f"Checking out to PR branch {branch_to_use} for issue {issue.number}"
                )

                if not branch_to_use:
                    raise ValueError("Branch name cannot be None")

                # Fetch the branch first to ensure it exists locally
                fetch_cmd = ["git", "fetch", "origin", branch_to_use]
                subprocess.check_output(
                    fetch_cmd,
                    cwd=repo_dir,
                )

                # Checkout the branch
                checkout_cmd = ["git", "checkout", branch_to_use]
                subprocess.check_output(
                    checkout_cmd,
                    cwd=repo_dir,
                )

                base_commit = (
                    subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=repo_dir)
                    .decode("utf-8")
                    .strip()
                )

            logger.info(
                f"Issue Resolve - Using base commit {base_commit} for issue {issue.number}\nOutput successfully written to {output_file}"
            )

            output = await self.process_issue(
                issue,
                base_commit,
                self.issue_handler,
                reset_logger,
            )

            customOutput = CustomResolverOutput(**output.model_dump())
            
            # Extract cost information from metrics if available
            self._calculate_and_set_costs(customOutput)

        finally:
            logger.info("Finished.")

            end_time = time.time()
            duration = end_time - start_time
            logger.info(f"Total time taken: {duration} seconds")

            if customOutput is not None:  # Check if customOutput was created
                customOutput.duration = duration
                # Log cost information
                if customOutput.accumulated_cost is not None:
                    logger.info(f"Total cost for issue resolution: ${customOutput.accumulated_cost:.4f}")
                # logger.info(f"Output: {customOutput}")
                return customOutput
            else:
                # Create a minimal error output if something went wrong
                error_output = CustomResolverOutput(
                    issue=issue,
                    issue_type=self.issue_type,
                    instruction="",
                    base_commit=base_commit,
                    git_patch=None,
                    history=[],
                    metrics={},
                    success=False,
                    comment_success=None,
                    result_explanation="Error occurred during processing",
                    error="Failed to complete processing",
                    duration=duration,
                )
                return error_output

    async def get_new_commit_hash(
        self, output_dir, resolver_output: CustomResolverOutput, pr_type: str
    ) -> None:
        # 1) initialize_repo
        patched_repo_dir = initialize_repo(
            output_dir=output_dir,
            issue_number=resolver_output.issue.number,
            issue_type=resolver_output.issue_type,
            base_commit=resolver_output.base_commit,
        )

        # logger.info(f"[DEBUG] Previous Patched Repo Dir: {patched_repo_dir}")
        branch_name, default_branch, base_url, headers = None, None, None, None

        if resolver_output.git_patch:
            # 2) apply_patch
            apply_patch(patched_repo_dir, resolver_output.git_patch)

            # 3) make_commit [NEW!] with Summary
            # logger.info(f"[DEBUG] Resolver Output: {resolver_output} to {output_dir}")
            # make_commit(patched_repo_dir, resolver_output.issue, resolver_output.issue_type)
            make_commit_with_summary(
                patched_repo_dir,
                resolver_output.issue,
                resolver_output.issue_type,
                resolver_output,
                branch_name,
                output_dir,
            )

            # 4) branch checkout and push
            branch_name, default_branch, base_url, headers = (
                self.prepare_branch_and_push(
                    patch_dir=patched_repo_dir,
                    pr_type=pr_type,
                )
            )

        resolver_output.branch_name = branch_name
        resolver_output.default_branch = default_branch
        resolver_output.base_url = base_url
        resolver_output.headers = headers

        # 5) Retrieve commit hash
        rev_parse_cmd = f'git -C "{patched_repo_dir}" rev-parse HEAD'
        result = subprocess.run(
            rev_parse_cmd, shell=True, capture_output=True, text=True
        )
        new_hash = result.stdout.strip()

        # 6) Assign it back to the resolver_output
        resolver_output.commit_hash = new_hash
        resolver_output.repo_dir = patched_repo_dir

        return

    async def _resolve_with_model(
        self, model_index: int, output_dir: str
    ) -> CustomResolverOutput:
        """Resolve issue with a specific model configuration.

        This method creates a completely independent resolver context for parallel execution,
        ensuring no shared state between concurrent model runs.

        Args:
            model_index: Index of the model configuration to use (0 or 1)
            output_dir: Output directory for this model's results

        Returns:
            CustomResolverOutput: The resolved output for this model
        """
        llm_config = self.llm_configs[model_index]

        # Create a completely independent issue handler factory for this model
        # This ensures thread safety and no shared state between parallel executions
        factory = IssueHandlerFactory(
            owner=self.owner,
            repo=self.repo,
            token=self.token,
            username=self.username,
            platform=self.platform,
            base_domain="github.com",
            issue_type=self.issue_type,
            llm_config=llm_config,
        )

        # Create independent issue handler for this execution
        issue_handler = factory.create()

        # Store current instance state to avoid conflicts
        current_llm_config = self.llm_config if hasattr(self, "llm_config") else None
        current_issue_handler = (
            self.issue_handler if hasattr(self, "issue_handler") else None
        )
        current_output_dir = self.output_dir if hasattr(self, "output_dir") else None
        current_error_tracker_issue_info = (
            self.error_tracker.issue_title,
            self.error_tracker.issue_body,
        )

        # Create a temporary backup of instance state for safe parallel execution
        backup_state = {
            "llm_config": current_llm_config,
            "issue_handler": current_issue_handler,
            "output_dir": current_output_dir,
        }

        try:
            # Temporarily set instance variables for this model's execution
            # This allows resolve_issue() to access the correct configuration
            self.llm_config = llm_config
            self.issue_handler = issue_handler
            self.output_dir = output_dir

            # Resolve the issue with this specific model configuration
            # This is the main time-consuming operation that benefits from parallelization
            resolver_output = await self.resolve_issue()

            # Update the output with the correct model name for identification
            resolver_output_dict = resolver_output.model_dump()
            resolver_output_dict["model"] = llm_config.model.split("/")[-1]
            resolved_output = CustomResolverOutput(**resolver_output_dict)
            
            # Calculate costs for this specific model's execution
            self._calculate_and_set_costs_for_model(resolved_output, llm_config)

            return resolved_output

        except Exception as e:
            # Ensure proper exception handling for parallel execution
            logger.error(
                f"Error in _resolve_with_model for model {model_index}: {str(e)}"
            )
            raise

        finally:
            # Restore the original instance state to avoid affecting other operations
            # This is critical for maintaining consistency in parallel execution
            for key, value in backup_state.items():
                if value is not None:
                    setattr(self, key, value)
                elif hasattr(self, key):
                    delattr(self, key)

            # Restore error tracker issue info if it was changed
            if current_error_tracker_issue_info != (None, None):
                self.error_tracker.set_issue_info(*current_error_tracker_issue_info)

    def prepare_branch_and_push(
        self,
        patch_dir: str,
        pr_type: str,
    ) -> tuple[str, str, str, dict]:
        """
        1) Validates pr_type.
        2) Sets up API headers, base_url.
        3) Generates a unique branch name.
        4) Gets the repository's default branch.
        5) Creates & checks out a new local branch.
        6) Pushes the new branch to GitHub.
        7) Returns the data needed for creating a PR (branch name, default branch, base_url, etc.)
        plus the commit hash of HEAD.
        """

        if pr_type not in ["branch", "draft", "ready"]:
            raise ValueError(f"Invalid pr_type: {pr_type}")

        # Prepare GitHub API details
        headers = {
            "Authorization": f"token {self.token}",
            "Accept": "application/vnd.github.v3+json",
        }
        base_url = f"https://api.github.com/repos/{self.owner}/{self.repo}"

        # Create a new branch name
        base_branch_name = f"openhands-fix-issue-{self.issue_number}"
        branch_name = base_branch_name + "-try1"
        attempt = 1

        # Ensure the branch doesn't already exist on the remote
        while (
            httpx.get(f"{base_url}/branches/{branch_name}", headers=headers).status_code
            == 200
        ):
            attempt += 1
            branch_name = f"{base_branch_name}-try{attempt}"

        # Get the default branch
        response = requests.get(base_url, headers=headers)
        response.raise_for_status()
        default_branch = response.json()["default_branch"]

        # Create and checkout the new branch locally
        result = subprocess.run(
            f"git -C {shlex.quote(patch_dir)} checkout -b {shlex.quote(branch_name)}",
            shell=True,
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            print(f"Error creating new branch: {result.stderr}")
            raise RuntimeError(
                f"Failed to create a new branch {branch_name} in {patch_dir}."
            )

        # Determine the repository to push to
        push_owner = self.owner
        push_repo = self.repo

        # Construct push command
        username_and_token = (
            f"{self.username}:{self.token}"
            if self.username
            else f"x-auth-token:{self.token}"
        )
        push_command = (
            f"git -C {shlex.quote(patch_dir)} push "
            f"https://{username_and_token}@github.com/"
            f"{push_owner}/{push_repo}.git {shlex.quote(branch_name)}"
        )
        result = subprocess.run(
            push_command, shell=True, capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"Error pushing changes\n{push_command}\n{result.stderr}")
            raise RuntimeError("Failed to push changes to the remote repository")

        return branch_name, default_branch, base_url, headers

    def _calculate_and_set_costs(self, resolver_output: CustomResolverOutput) -> None:
        """Calculate and set cost information from OpenHands metrics."""
        try:
            # Extract metrics from the resolver output
            if resolver_output.metrics and hasattr(resolver_output.metrics, 'accumulated_cost'):
                resolver_output.accumulated_cost = resolver_output.metrics.accumulated_cost
            elif resolver_output.metrics and isinstance(resolver_output.metrics, dict):
                resolver_output.accumulated_cost = resolver_output.metrics.get('accumulated_cost', 0.0)
            
            # Extract token usage information
            if resolver_output.metrics and hasattr(resolver_output.metrics, 'accumulated_token_usage'):
                resolver_output.token_usage = resolver_output.metrics.accumulated_token_usage
            elif resolver_output.metrics and isinstance(resolver_output.metrics, dict):
                resolver_output.token_usage = resolver_output.metrics.get('accumulated_token_usage', {})
            
            # Get cost per token from LLM config if available
            if hasattr(self, 'llm_config') and self.llm_config:
                resolver_output.cost_per_input_token = self.llm_config.input_cost_per_token
                resolver_output.cost_per_output_token = self.llm_config.output_cost_per_token
                
                # If accumulated cost isn't available but we have token counts and rates, calculate it
                if not resolver_output.accumulated_cost and resolver_output.token_usage:
                    input_tokens = 0
                    output_tokens = 0
                    
                    if isinstance(resolver_output.token_usage, dict):
                        # Sum up tokens from different components
                        for component_usage in resolver_output.token_usage.values():
                            if isinstance(component_usage, dict):
                                input_tokens += component_usage.get('prompt_tokens', 0)
                                output_tokens += component_usage.get('completion_tokens', 0)
                    
                    # Calculate cost
                    input_cost = input_tokens * (self.llm_config.input_cost_per_token or 0)
                    output_cost = output_tokens * (self.llm_config.output_cost_per_token or 0)
                    resolver_output.accumulated_cost = input_cost + output_cost
                    
        except Exception as e:
            logger.warning(f"Failed to calculate costs for resolver output: {e}")
            resolver_output.accumulated_cost = 0.0
            resolver_output.token_usage = {}

    def _calculate_and_set_costs_for_model(self, resolver_output: CustomResolverOutput, llm_config: "LLMConfig") -> None:
        """Calculate and set cost information for a specific model configuration."""
        try:
            # Extract metrics from the resolver output
            if resolver_output.metrics and hasattr(resolver_output.metrics, 'accumulated_cost'):
                resolver_output.accumulated_cost = resolver_output.metrics.accumulated_cost
            elif resolver_output.metrics and isinstance(resolver_output.metrics, dict):
                resolver_output.accumulated_cost = resolver_output.metrics.get('accumulated_cost', 0.0)
            
            # Extract token usage information
            if resolver_output.metrics and hasattr(resolver_output.metrics, 'accumulated_token_usage'):
                resolver_output.token_usage = resolver_output.metrics.accumulated_token_usage
            elif resolver_output.metrics and isinstance(resolver_output.metrics, dict):
                resolver_output.token_usage = resolver_output.metrics.get('accumulated_token_usage', {})
            
            # Set cost per token from the specific LLM config
            resolver_output.cost_per_input_token = llm_config.input_cost_per_token
            resolver_output.cost_per_output_token = llm_config.output_cost_per_token
            
            # If accumulated cost isn't available but we have token counts and rates, calculate it
            if not resolver_output.accumulated_cost and resolver_output.token_usage:
                input_tokens = 0
                output_tokens = 0
                
                if isinstance(resolver_output.token_usage, dict):
                    # Sum up tokens from different components
                    for component_usage in resolver_output.token_usage.values():
                        if isinstance(component_usage, dict):
                            input_tokens += component_usage.get('prompt_tokens', 0)
                            output_tokens += component_usage.get('completion_tokens', 0)
                
                # Calculate cost
                input_cost = input_tokens * (llm_config.input_cost_per_token or 0)
                output_cost = output_tokens * (llm_config.output_cost_per_token or 0)
                resolver_output.accumulated_cost = input_cost + output_cost
                
            logger.info(f"Cost calculation for {llm_config.model}: ${resolver_output.accumulated_cost:.4f}")
                
        except Exception as e:
            logger.warning(f"Failed to calculate costs for model {llm_config.model}: {e}")
            resolver_output.accumulated_cost = 0.0
            resolver_output.token_usage = {}


def main() -> None:
    import argparse

    def int_or_none(value: str) -> int | None:
        if value.lower() == "none":
            return None
        else:
            return int(value)

    parser = argparse.ArgumentParser(description="Resolve issues from Github.")
    parser.add_argument(
        "--selected-repo",
        type=str,
        required=True,
        help="repository to resolve issues in form of `owner/repo`.",
    )
    parser.add_argument(
        "--token",
        type=str,
        default=None,
        help="token to access the repository.",
    )
    parser.add_argument(
        "--username",
        type=str,
        default=None,
        help="username to access the repository.",
    )
    parser.add_argument(
        "--base-container-image",
        type=str,
        default=None,
        help="base container image to use.",
    )
    parser.add_argument(
        "--runtime-container-image",
        type=str,
        default=None,
        help="Container image to use.",
    )
    parser.add_argument(
        "--max-iterations",
        type=int,
        default=50,
        help="Maximum number of iterations to run.",
    )
    parser.add_argument(
        "--issue-number",
        type=int,
        required=True,
        help="Issue number to resolve.",
    )
    parser.add_argument(
        "--comment-id",
        type=int_or_none,
        required=False,
        default=None,
        help="Resolve a specific comment",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="output",
        help="Output directory to write the results.",
    )
    parser.add_argument(
        "--llm-model",
        type=str,
        default="Mock GPT",
        help="Mock model name to adapt with the existing code.",
    )
    parser.add_argument(
        "--llm-models",
        type=str,
        default=None,
        help="LLM models to use.",
    )
    parser.add_argument(
        "--llm-api-key",
        type=str,
        default=None,
        help="LLM API key to use.",
    )
    parser.add_argument(
        "--llm-base-url",
        type=str,
        default=None,
        help="LLM base URL to use.",
    )
    parser.add_argument(
        "--prompt-file",
        type=str,
        default=None,
        help="Path to the prompt template file in Jinja format.",
    )
    parser.add_argument(
        "--repo-instruction-file",
        type=str,
        default=None,
        help="Path to the repository instruction file in text format.",
    )
    parser.add_argument(
        "--issue-type",
        type=str,
        default="issue",
        choices=["issue", "pr"],
        help="Type of issue to resolve, either open issue or pr comments.",
    )
    parser.add_argument(
        "--is-experimental",
        type=lambda x: x.lower() == "true",
        help="Whether to run in experimental mode.",
    )
    parser.add_argument(
        "--base-domain",
        type=str,
        default=None,
        help='Base domain for the git server (defaults to "github.com" for GitHub and "gitlab.com" for GitLab)',
    )

    my_args = parser.parse_args()
    issue_resolver = PRArenaIssueResolver(my_args)

    asyncio.run(issue_resolver.resolve_issues_with_random_models())


if __name__ == "__main__":
    main()
