from __future__ import annotations

import hashlib
import json
import os
import shutil
from importlib import resources
from pathlib import Path
from typing import Mapping, Optional

DEFAULT_DATA_ROOT = "/root/data"


def data_root_from_env(env: Optional[Mapping[str, str]] = None) -> Path:
    source = env or os.environ
    return Path(source.get("DATA_ROOT", DEFAULT_DATA_ROOT)).expanduser()


def runtime_base_dir(env: Optional[Mapping[str, str]] = None) -> Path:
    return data_root_from_env(env) / ".lehome-cli" / "runtime"


def runtime_current_dir(env: Optional[Mapping[str, str]] = None) -> Path:
    return runtime_base_dir(env) / "current"


def bundle_root():
    return resources.files("lehome_cli.bundle")


def packaged_manifest() -> dict:
    manifest_path = bundle_root().joinpath("manifest.json")
    return json.loads(manifest_path.read_text(encoding="utf-8"))


def current_manifest(env: Optional[Mapping[str, str]] = None) -> Optional[dict]:
    manifest_path = runtime_current_dir(env) / "manifest.json"
    if not manifest_path.is_file():
        return None
    return json.loads(manifest_path.read_text(encoding="utf-8"))


def sha256_path(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def is_work_repo_root(path: Path) -> bool:
    return (
        path.is_dir()
        and (path / "pyproject.toml").is_file()
        and (path / "configs").is_dir()
        and (path / "scripts" / "eval.py").is_file()
    )


def find_work_repo(
    env: Optional[Mapping[str, str]] = None,
    cwd: Optional[Path] = None,
) -> Optional[Path]:
    source = env or os.environ
    dirname = source.get("LEHOME_DEFAULT_DIRNAME", "lehome-challenge")
    work_cwd = (cwd or Path.cwd()).resolve()
    candidates = [
        work_cwd,
        work_cwd / dirname,
        work_cwd.parent / dirname,
        data_root_from_env(source) / dirname,
    ]

    seen: set[Path] = set()
    for candidate in candidates:
        resolved = candidate.resolve()
        if resolved in seen:
            continue
        seen.add(resolved)
        if is_work_repo_root(resolved):
            return resolved
    return None


def _copy_resource_tree(source, destination: Path) -> None:
    if source.is_dir():
        destination.mkdir(parents=True, exist_ok=True)
        for child in source.iterdir():
            _copy_resource_tree(child, destination / child.name)
        return

    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_bytes(source.read_bytes())


def _mark_shells_executable(root: Path) -> None:
    for shell_file in root.rglob("*.sh"):
        shell_file.chmod(shell_file.stat().st_mode | 0o755)


def _sync_runtime_links(env: Mapping[str, str], current_dir: Path) -> None:
    data_root = data_root_from_env(env)
    dirname = env.get("LEHOME_DEFAULT_DIRNAME", "lehome-challenge")
    targets = {dirname, "lehome-challenge"}

    for name in targets:
        link_path = current_dir / name
        target_path = data_root / dirname
        if link_path.exists() or link_path.is_symlink():
            if link_path.is_symlink() and link_path.resolve(strict=False) == target_path.resolve(strict=False):
                continue
            if link_path.is_dir() and not link_path.is_symlink():
                continue
            link_path.unlink()
        link_path.symlink_to(target_path)


def _manifest_matches(root: Path, manifest: dict) -> bool:
    for entry in manifest.get("files", []):
        file_path = root / entry["bundle_relpath"]
        if not file_path.is_file():
            return False
        if sha256_path(file_path) != entry["sha256"]:
            return False
    return True


def runtime_is_current(env: Optional[Mapping[str, str]] = None) -> bool:
    current_root = runtime_current_dir(env)
    manifest = packaged_manifest()
    active_manifest = current_manifest(env)
    if active_manifest is None:
        return False
    if active_manifest.get("bundle_version") != manifest.get("bundle_version"):
        return False
    return _manifest_matches(current_root, manifest)


def extract_bundle(env: Optional[Mapping[str, str]] = None) -> Path:
    active_env = dict(env or os.environ)
    base_dir = runtime_base_dir(env)
    current_dir = runtime_current_dir(env)
    staging_dir = base_dir / f".staging-{os.getpid()}"
    backup_dir = base_dir / f".backup-{os.getpid()}"

    base_dir.mkdir(parents=True, exist_ok=True)
    if staging_dir.exists():
        shutil.rmtree(staging_dir)
    if backup_dir.exists():
        shutil.rmtree(backup_dir)

    _copy_resource_tree(bundle_root(), staging_dir)
    _mark_shells_executable(staging_dir)
    _sync_runtime_links(active_env, staging_dir)

    if current_dir.exists():
        current_dir.rename(backup_dir)
    staging_dir.rename(current_dir)
    if backup_dir.exists():
        shutil.rmtree(backup_dir)
    return current_dir


def ensure_runtime(env: Optional[Mapping[str, str]] = None) -> Path:
    active_env = dict(env or os.environ)
    if runtime_is_current(active_env):
        current_dir = runtime_current_dir(active_env)
        _sync_runtime_links(active_env, current_dir)
        return current_dir
    return extract_bundle(active_env)
