"""Unidiff extraction with before/after context."""

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import git
from git import Repo


@dataclass
class DiffBundle:
    repo: str
    commit_sha: str
    diff: str
    before_files: dict
    after_files: dict


def _detect_language(filename: str) -> str:
    ext = Path(filename).suffix.lower()
    mapping = {
        ".py": "python",
        ".java": "java",
        ".js": "javascript",
        ".ts": "javascript",
        ".jsx": "javascript",
        ".tsx": "javascript",
    }
    return mapping.get(ext, "unknown")


def extract_diff_bundle(
    repo_path: str,
    commit_sha: str,
    parent_index: int = 0,
) -> DiffBundle:
    """Extract unified diff and before/after file snapshots for a commit."""
    repo = Repo(str(repo_path))
    commit = repo.commit(commit_sha)
    parents = commit.parents
    if parents:
        base = parents[parent_index]
        diff = base.diff(commit, create_patch=True, unified=3)
    else:
        # Initial commit: diff against empty tree
        diff = commit.diff(git.Git(repo).hash_object("/dev/null", t=None), create_patch=True, unified=3)

    diff_text = "\n".join(d.diff.decode("utf-8", errors="replace") for d in diff if d.diff)

    before_files = {}
    after_files = {}
    for d in diff:
        a_path = d.a_path or d.b_path
        b_path = d.b_path or d.a_path
        if a_path and d.a_blob:
            try:
                before_files[a_path] = d.a_blob.data_stream.read().decode("utf-8", errors="replace")
            except Exception:
                before_files[a_path] = ""
        if b_path and d.b_blob:
            try:
                after_files[b_path] = d.b_blob.data_stream.read().decode("utf-8", errors="replace")
            except Exception:
                after_files[b_path] = ""

    return DiffBundle(
        repo=Path(repo_path).name,
        commit_sha=commit_sha,
        diff=diff_text,
        before_files=before_files,
        after_files=after_files,
    )


def bundle_to_sample_dict(bundle: DiffBundle, primary_language: Optional[str] = None) -> dict:
    """Convert a DiffBundle into a dict matching the Sample schema."""
    if not primary_language:
        # Infer from changed files
        exts = {Path(p).suffix.lower() for p in bundle.after_files or bundle.before_files}
        for ext, lang in {".py": "python", ".java": "java", ".js": "javascript"}.items():
            if ext in exts:
                primary_language = lang
                break
        primary_language = primary_language or "unknown"

    return {
        "repo": bundle.repo,
        "commit_sha": bundle.commit_sha,
        "language": primary_language,
        "diff": bundle.diff,
        "before_context": bundle.before_files,
        "after_context": bundle.after_files,
    }