"""PR creation and description helpers for performance results."""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING

from codeflash_core import (
    FileDiffContent,
    PlatformClient,
    PrComment,
    check_and_push_branch,
    get_repo_owner_and_name,
    performance_gain,
)

from ..ai._tabulate import tabulate
from ..testing._testgen import format_perf, format_time

if TYPE_CHECKING:
    import git

    from ..test_discovery.models import FunctionCalledInTest
    from ..testing.models import InvocationId, TestConfig, TestFiles

log = logging.getLogger(__name__)


def _get_pr_number() -> int | None:
    """Return the PR number from the environment, or *None*."""
    raw = os.environ.get("CODEFLASH_PR_NUMBER")
    if raw:
        return int(raw)
    return None


def _github_pr_url(owner: str, repo: str, pr_id: str) -> str:
    """Build a GitHub PR URL from components."""
    return f"https://github.com/{owner}/{repo}/pull/{pr_id}"


def check_create_pr(  # noqa: PLR0913
    *,
    platform_client: PlatformClient,
    git_repo: git.Repo,
    original_code: dict[str, str],
    new_code: dict[str, str],
    pr_comment: PrComment,
    existing_tests: str,
    generated_tests: str,
    function_trace_id: str,
    coverage_message: str,
    replay_tests: str = "",
    concolic_tests: str = "",
    optimization_review: str = "",
    git_remote: str = "origin",
) -> None:
    """Create or update a PR with the optimized code.

    When ``CODEFLASH_PR_NUMBER`` is set, suggests changes to
    the existing PR.  Otherwise creates a new PR via the
    platform backend.
    """
    pr_number = _get_pr_number()

    file_changes: dict[str, FileDiffContent] = {
        p: FileDiffContent(
            old_content=original_code[p],
            new_content=new_code[p],
        )
        for p in original_code
        if original_code[p] != new_code[p]
    }
    if not file_changes:
        log.info("No file changes to submit for PR.")
        return

    if pr_number is not None:
        log.info(
            "Suggesting changes to PR #%d …",
            pr_number,
        )
        owner, repo = get_repo_owner_and_name(
            git_repo,
            git_remote,
        )
        response = platform_client.suggest_changes(
            owner=owner,
            repo=repo,
            pr_number=pr_number,
            file_changes=file_changes,
            pr_comment=pr_comment,
            existing_tests=existing_tests,
            generated_tests=generated_tests,
            trace_id=function_trace_id,
            coverage_message=coverage_message,
            replay_tests=replay_tests,
            concolic_tests=concolic_tests,
            optimization_review=optimization_review,
        )
        if response.ok:
            log.info(
                "Suggestions were successfully made to PR #%d",
                pr_number,
            )
        else:
            log.error(
                "Failed to suggest changes to PR #%d: %s",
                pr_number,
                response.text,
            )
    else:
        log.info("Creating a new PR with the optimized code…")
        owner, repo = get_repo_owner_and_name(
            git_repo,
            git_remote,
        )
        if not check_and_push_branch(
            git_repo,
            git_remote,
            wait_for_push=True,
        ):
            log.warning(
                "Branch is not pushed, skipping PR creation…",
            )
            return
        base_branch = git_repo.active_branch.name
        response = platform_client.create_pr(
            owner=owner,
            repo=repo,
            base_branch=base_branch,
            file_changes=file_changes,
            pr_comment=pr_comment,
            existing_tests=existing_tests,
            generated_tests=generated_tests,
            trace_id=function_trace_id,
            coverage_message=coverage_message,
            replay_tests=replay_tests,
            concolic_tests=concolic_tests,
            optimization_review=optimization_review,
        )
        if response.ok:
            pr_id = response.text
            pr_url = _github_pr_url(owner, repo, pr_id)
            log.info(
                "Created PR #%s: %s",
                pr_id,
                pr_url,
            )
        elif response.status_code == 401 and "not installed" in response.text:  # noqa: PLR2004
            log.warning(
                "Could not create PR: the Codeflash GitHub App is not installed on %s/%s. "
                "Install it at https://github.com/apps/codeflash-ai to enable automatic PRs.",
                owner,
                repo,
            )
        else:
            log.error(
                "Failed to create PR: %s",
                response.text,
            )


def existing_tests_source_for(  # noqa: C901, PLR0912, PLR0913, PLR0915
    function_qualified_name_with_modules_from_root: str,
    function_to_tests: dict[str, set[FunctionCalledInTest]],
    test_cfg: TestConfig,
    original_runtimes_all: dict[InvocationId, list[int]],
    optimized_runtimes_all: dict[InvocationId, list[int]],
    test_files_registry: TestFiles | None = None,
) -> tuple[str, str, str]:
    """Build markdown tables summarising existing-test performance.

    Returns a triple of markdown strings:
    ``(existing_table, replay_table, concolic_table)``.
    """
    test_files = function_to_tests.get(
        function_qualified_name_with_modules_from_root,
    )
    if not test_files:
        return "", "", ""

    rows_existing: list[list[str]] = []
    rows_concolic: list[list[str]] = []
    rows_replay: list[list[str]] = []
    headers = [
        "Test File::Test Function",
        "Original \u23f1\ufe0f",
        "Optimized \u23f1\ufe0f",
        "Speedup",
    ]
    tests_root = test_cfg.tests_root

    original_tests_to_runtimes: dict[Path, dict[str, int]] = {}
    optimized_tests_to_runtimes: dict[Path, dict[str, int]] = {}

    # Build lookup from instrumented path -> original path
    instrumented_to_original: dict[Path, Path] = {}
    if test_files_registry:
        for registry_tf in test_files_registry.test_files:
            if registry_tf.original_file_path:
                if registry_tf.instrumented_behavior_file_path:
                    instrumented_to_original[
                        registry_tf.instrumented_behavior_file_path.resolve()
                    ] = registry_tf.original_file_path.resolve()
                if registry_tf.benchmarking_file_path:
                    instrumented_to_original[
                        registry_tf.benchmarking_file_path.resolve()
                    ] = registry_tf.original_file_path.resolve()

    # Resolve all paths to absolute for consistent comparison
    non_generated_tests: set[Path] = set()
    for test_file in test_files:
        resolved = test_file.tests_in_file.test_file.resolve()
        non_generated_tests.add(resolved)

    all_invocation_ids = (
        original_runtimes_all.keys() | optimized_runtimes_all.keys()
    )
    for invocation_id in all_invocation_ids:
        test_module_path = invocation_id.test_module_path
        # Python: convert module name to path
        abs_path = (
            Path(test_module_path.replace(".", os.sep))
            .with_suffix(".py")
            .resolve()
        )

        if abs_path not in non_generated_tests:
            continue

        if abs_path not in original_tests_to_runtimes:
            original_tests_to_runtimes[abs_path] = {}
        if abs_path not in optimized_tests_to_runtimes:
            optimized_tests_to_runtimes[abs_path] = {}

        cls = invocation_id.test_class_name or ""
        func = invocation_id.test_function_name or ""
        qualified_name: str = f"{cls}.{func}" if cls else func

        if qualified_name not in original_tests_to_runtimes[abs_path]:
            original_tests_to_runtimes[abs_path][qualified_name] = 0
        if qualified_name not in optimized_tests_to_runtimes[abs_path]:
            optimized_tests_to_runtimes[abs_path][qualified_name] = 0
        if invocation_id in original_runtimes_all:
            original_tests_to_runtimes[abs_path][qualified_name] += min(
                original_runtimes_all[invocation_id],
            )
        if invocation_id in optimized_runtimes_all:
            optimized_tests_to_runtimes[abs_path][qualified_name] += min(
                optimized_runtimes_all[invocation_id],
            )

    # Build result tables
    all_abs_paths = original_tests_to_runtimes.keys()
    for filename in sorted(all_abs_paths):
        all_qualified_names = original_tests_to_runtimes[filename].keys()
        for qualified_name in sorted(all_qualified_names):
            orig_rt = original_tests_to_runtimes[filename][qualified_name]
            opt_rt = optimized_tests_to_runtimes[filename][qualified_name]
            if orig_rt == 0 or opt_rt == 0:
                continue

            print_optimized_runtime = format_time(opt_rt)
            print_original_runtime = format_time(orig_rt)
            print_filename = (
                filename.resolve()
                .relative_to(Path(tests_root).resolve())
                .as_posix()
            )
            greater = opt_rt > orig_rt
            perf_gain = format_perf(
                performance_gain(
                    original_runtime_ns=orig_rt,
                    optimized_runtime_ns=opt_rt,
                )
                * 100,
            )

            row = [
                f"`{print_filename}::{qualified_name}`",
                f"{print_original_runtime}",
                f"{print_optimized_runtime}",
            ]

            if greater:
                row.append(f"{perf_gain}%\u26a0\ufe0f")
            else:
                row.append(f"{perf_gain}%\u2705")

            if "__replay_test_" in str(print_filename):
                rows_replay.append(row)
            elif "codeflash_concolic" in str(print_filename):
                rows_concolic.append(row)
            else:
                rows_existing.append(row)

    output_existing = tabulate(
        headers=headers,
        tabular_data=rows_existing,
        tablefmt="pipe",
        colglobalalign=None,
        preserve_whitespace=True,
    )
    output_existing += "\n"
    if len(rows_existing) == 0:
        output_existing = ""

    output_concolic = tabulate(
        headers=headers,
        tabular_data=rows_concolic,
        tablefmt="pipe",
        colglobalalign=None,
        preserve_whitespace=True,
    )
    output_concolic += "\n"
    if len(rows_concolic) == 0:
        output_concolic = ""

    output_replay = tabulate(
        headers=headers,
        tabular_data=rows_replay,
        tablefmt="pipe",
        colglobalalign=None,
        preserve_whitespace=True,
    )
    output_replay += "\n"
    if len(rows_replay) == 0:
        output_replay = ""

    return output_existing, output_replay, output_concolic
