from datetime import datetime, UTC
import re
from time import time
from typing import Literal, TypedDict
from collections.abc import Iterable

from drafteleu.bump import bump_version
from drafteleu.changelog import compile_changelog, get_prs, PullRequest
from drafteleu.releases import (
    find_latest_release,
    create_or_update_draft_release,
    Release,
)
from drafteleu.async_utils import create_task_id_context, get_async_logger
import asyncio

logger = get_async_logger(__name__)


class Configuration(TypedDict, total=False):
    include_labels: list
    tag_regex: str
    stage: Literal["pre", "final", "both"]
    version_template: str
    tag_template: str
    title_template: str
    rc_template: str
    bump_part: Literal["major", "minor", "patch"]


async def preload_prs(
    repo_slug: str, configurations: Iterable[Configuration]
) -> list[PullRequest]:
    start = time()

    all_labels = set()
    earliest_release_date = datetime.now(UTC)
    for config in configurations:
        include_labels = config.get("include_labels", [])
        all_labels.update(include_labels)
        tag_regex = config.get("tag_regex", "v(?P<version>.+)")
        stage = config.get("stage", "both")

        for current_stage in ("pre", "final") if stage == "both" else (stage,):
            include_prerelease = current_stage == "pre"
            latest_release = await find_latest_release(
                tag_regex,
                repo=repo_slug,
                include_prereleases=include_prerelease,
            )
            if latest_release is None:
                logger.info("No previous release found, assuming initial release")
                earliest_release_date = datetime.fromisoformat("1970-01-01T00:00:00Z")
                # No point in looking for earlier releases
                break
            else:
                release_date = latest_release.published_at
            if release_date < earliest_release_date:
                earliest_release_date = release_date

    # Now that we have the earliest release date and all labels, fetch PRs once
    prs = await get_prs(
        merged_since=earliest_release_date,
        repo=repo_slug,
        labels=list(all_labels),
    )
    logger.info(
        f"Preloaded {len(prs)} PRs since {earliest_release_date} in {time() - start:.2f}s"
    )
    return prs


def pr_has_any_label(pr: PullRequest, labels: set) -> bool:
    return not labels or labels.intersection(pr.labels)


async def process_config(
    config: Configuration,
    all_prs: list[PullRequest],
    repo_slug: str,
    stage: Literal["pre", "final", "both"],
    dry_run: bool,
):
    version_template = config.get("version_template", "{major}.{minor}.{patch}")
    tag_template = config.get("tag_template", "v{next_version}")
    title_template = config.get("title_template", "Release {next_version}")
    tag_regex = config.get("tag_regex", "v(?P<version>.+)")
    rc_template = config.get("rc_template", "-{pre_label}{pre_number}")
    bump_part = config.get("bump_part", "minor")
    include_labels = set(config.get("include_labels", []))

    for current_stage in ("pre", "final") if stage == "both" else (stage,):
        logger.info(f"Processing stage: {current_stage}")
        include_prerelease = current_stage == "pre"

        latest_release = await find_latest_release(
            tag_regex, repo=repo_slug, include_prereleases=include_prerelease
        )
        if latest_release is None:
            logger.info("No previous release found, assuming initial release")
            initial_version = version_template.format(major=0, minor=0, patch=0)
            latest_release = Release(
                tag_name=tag_template.format(next_version=initial_version),
                published_at=datetime.fromisoformat("1970-01-01T00:00:00Z"),
                is_draft=False,
            )
        else:
            logger.info(
                f"Latest release: {latest_release.tag_name} at {latest_release.published_at}"
            )

        prs = [
            pr
            for pr in all_prs
            if pr_has_any_label(pr, include_labels)
            and pr.closed_at > latest_release.published_at
        ]
        logger.info(
            f"Found {len(prs)} PRs since last release with labels {include_labels}"
        )

        changelog = compile_changelog(prs)
        logger.debug("Changelog:\n" + changelog)

        current_version = re.match(tag_regex, latest_release.tag_name).group("version")
        logger.info(f"Current version: {current_version}")
        next_version = bump_version(
            current_version,
            "dev" if include_prerelease else bump_part,
            version_template=version_template,
            rc_template=rc_template,
            rc_bump_part=bump_part,
        )
        next_tag = tag_template.format(next_version=next_version)
        next_title = title_template.format(
            next_version=next_version,
            next_tag=next_tag,
            current_version=current_version,
            current_tag=latest_release.tag_name,
            current_date=latest_release.published_at,
            stage=current_stage,
        )
        logger.info(f"Next version: {next_version}")
        logger.info(f"Next tag: {next_tag}")
        logger.info(f"Next title: {next_title}")

        await create_or_update_draft_release(
            tag_regex,
            next_tag,
            next_title,
            changelog,
            repo=repo_slug,
            dry_run=dry_run,
            prerelease=include_prerelease,
        )


async def run_drafteleu(
    repo_slug: str,
    *,
    stage: Literal["pre", "final", "both"] = "both",
    dry_run: bool = False,
    configurations: Iterable[Configuration],
):
    all_prs = await preload_prs(repo_slug, configurations)

    try:
        async with asyncio.TaskGroup() as tg:
            for index, config in enumerate(configurations):
                # set the task id in the context so that it can be used in the task
                # function and any child functions without having to explicitely
                # pass it
                tg.create_task(
                    process_config(
                        config,
                        all_prs=all_prs,
                        repo_slug=repo_slug,
                        stage=stage,
                        dry_run=dry_run,
                    ),
                    context=create_task_id_context(index),
                )
    except BaseException:
        logger.flush_async_logs(prepend_task_id=True)
        raise
    else:
        logger.flush_async_logs()
