from __future__ import annotations

from app.models.enums import SourceType
from app.retrieval.models import (
    CodeSnippet,
    DocExcerpt,
    GroupedRetrievalContext,
    HydratedHit,
    HydrationResult,
    RelatedCommit,
)


def group_hydrated_hits(result: HydrationResult) -> GroupedRetrievalContext:
    """Split hydrated hits into code_snippets, doc_excerpts, and related_commits."""
    deduped = _dedupe_hits(result.hits)
    code_snippets: list[CodeSnippet] = []
    doc_excerpts: list[DocExcerpt] = []
    related_commits: list[RelatedCommit] = []

    for hit in deduped:
        if hit.source_type == SourceType.CODE.value:
            if hit.file_path is None or hit.symbol_name is None:
                continue
            if hit.start_line is None or hit.end_line is None:
                continue
            code_snippets.append(
                CodeSnippet(
                    record_id=hit.record_id,
                    file_path=hit.file_path,
                    symbol_name=hit.symbol_name,
                    symbol_type=hit.symbol_type,
                    language=hit.language,
                    start_line=hit.start_line,
                    end_line=hit.end_line,
                    text=hit.text,
                    score=hit.score,
                    graph_expanded=hit.graph_expanded,
                )
            )
        elif hit.source_type == SourceType.DOCS.value:
            if hit.doc_path is None or hit.section_title is None:
                continue
            doc_excerpts.append(
                DocExcerpt(
                    record_id=hit.record_id,
                    doc_path=hit.doc_path,
                    section_title=hit.section_title,
                    section_level=hit.section_level,
                    text=hit.text,
                    score=hit.score,
                    graph_expanded=hit.graph_expanded,
                )
            )
        elif hit.source_type == SourceType.COMMIT.value:
            if hit.commit_sha is None:
                continue
            related_commits.append(
                RelatedCommit(
                    record_id=hit.record_id,
                    commit_sha=hit.commit_sha,
                    summary=hit.text,
                    impacted_symbols=hit.impacted_symbols,
                    changed_files=hit.changed_files,
                    score=hit.score,
                    author_name=hit.author_name,
                    author_email=hit.author_email,
                    committer_name=hit.committer_name,
                    committer_email=hit.committer_email,
                    authored_at=hit.authored_at,
                    committed_at=hit.committed_at,
                    commit_message=hit.commit_message,
                )
            )

    return GroupedRetrievalContext(
        hits=deduped,
        code_snippets=code_snippets,
        doc_excerpts=doc_excerpts,
        related_commits=related_commits,
        retrieval_count=result.retrieval_count,
        hydrated_count=result.hydrated_count,
        skipped_count=result.skipped_count,
    )


def _dedupe_hits(hits: list[HydratedHit]) -> list[HydratedHit]:
    best: dict[str, HydratedHit] = {}
    for hit in hits:
        existing = best.get(hit.record_id)
        if existing is None or hit.score > existing.score:
            best[hit.record_id] = hit
    return sorted(best.values(), key=lambda item: item.score, reverse=True)