from __future__ import annotations
import base64
import hashlib
import io
import json
import os
import tempfile
import tokenize
from pathlib import Path
from typing import Any
import csscompressor
from bs4 import BeautifulSoup
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519
from jsmin import jsmin
from ._exceptions import IntegrityError, MetadataSignatureError, NotJson
from .jsonValidator import is_valid_json_file
from .logger import get_logger
logger = get_logger(name="Integrity")
CODE_EXTENSIONS = {
    ".py", ".js", ".mjs", ".cjs", ".ts", ".tsx", ".go", ".java",
    ".c", ".h", ".cpp", ".cc", ".cxx", ".hpp", ".hh", ".hxx", ".cs",
    ".html", ".htm", ".css",
}
def _atomic_write_json(path: str | Path, data: Any) -> None:
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    temp_name: str | None = None
    try:
        with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=target.parent, delete=False) as tmp:
            temp_name = tmp.name
            json.dump(data, tmp, indent=4, sort_keys=True)
            tmp.flush()
            os.fsync(tmp.fileno())
        os.replace(temp_name, target)
    finally:
        if temp_name and os.path.exists(temp_name):
            try:
                os.remove(temp_name)
            except OSError:
                pass
def _clean_python(file_path: str) -> tuple[bytes | None, int]:
    clean_code = io.StringIO()
    line_count = 0
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            for token_type, token_string, start, end, line in tokenize.generate_tokens(f.readline):
                if token_type == tokenize.COMMENT:
                    continue
                if token_type == tokenize.STRING and start[1] == 0:
                    continue
                clean_code.write(token_string)
                if token_type in (tokenize.NEWLINE, tokenize.NL):
                    line_count += 1
        return clean_code.getvalue().encode("utf-8"), line_count
    except Exception as e:
        logger.debug(f"Failed to clean python code for {file_path}: {e}")
        return None, 0
def _clean_js_like(src: str) -> str:
    try:
        return jsmin(src)
    except Exception as e:
        logger.debug(f"jsmin failed: {e}")
        return src
def _clean_css(src: str) -> str:
    try:
        return csscompressor.compress(src)
    except Exception as e:
        logger.debug(f"csscompressor failed: {e}")
        return src
def _clean_html(src: str) -> str:
    try:
        soup = BeautifulSoup(src, "html.parser")
        return soup.prettify()
    except Exception as e:
        logger.debug(f"BeautifulSoup failed: {e}")
        return src
def get_clean_code(file_path: str) -> tuple[bytes | None, int]:
    _, ext = os.path.splitext(file_path.lower())
    if ext == ".py":
        return _clean_python(file_path)
    try:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            src = f.read()
    except Exception as e:
        logger.debug(f"Failed to read source for clean {file_path}: {e}")
        return None, 0
    if ext in {".js", ".mjs", ".cjs", ".ts", ".tsx"}:
        cleaned = _clean_js_like(src)
        return cleaned.encode("utf-8"), cleaned.count("\\n") + 1 if cleaned else 0
    if ext == ".css":
        cleaned = _clean_css(src)
        return cleaned.encode("utf-8"), cleaned.count("\\n") + 1 if cleaned else 0
    if ext in {".html", ".htm"}:
        cleaned = _clean_html(src)
        return cleaned.encode("utf-8"), cleaned.count("\\n") + 1 if cleaned else 0
    return None, 0
def calculate_hash(file_path: str, *, algorithm: str = "sha256") -> tuple[str, int]:
    try:
        hasher = hashlib.new(algorithm)
    except ValueError as e:
        raise ValueError(f"Unsupported hash algorithm: {algorithm}") from e
    _, ext = os.path.splitext(file_path.lower())
    if ext in CODE_EXTENSIONS:
        code_bytes, lines = get_clean_code(file_path)
        if code_bytes is not None:
            hasher.update(code_bytes)
            return hasher.hexdigest(), lines
    line_count = 0
    with open(file_path, "rb") as f:
        for byte_block in iter(lambda: f.read(65536), b""):
            hasher.update(byte_block)
    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
        line_count = sum(1 for _ in f)
    return hasher.hexdigest(), line_count
def generate_keys(private_key_path: str) -> str:
    private_key = ed25519.Ed25519PrivateKey.generate()
    public_key = private_key.public_key()
    path = Path(private_key_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "wb") as f:
        f.write(
            private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption(),
            )
        )
    public_bytes = public_key.public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    public_b64 = base64.b64encode(public_bytes).decode("ascii")
    logger.info(f"Private key saved to {path}")
    logger.info("Store the public key securely")
    logger.info(public_b64)
    return public_b64
def _load_private_key(private_key_path: str) -> ed25519.Ed25519PrivateKey:
    path = Path(private_key_path)
    if not path.exists():
        raise FileNotFoundError(f"Private key file not found at {path}")
    with open(path, "rb") as f:
        private_key = serialization.load_pem_private_key(f.read(), password=None)
    if not isinstance(private_key, ed25519.Ed25519PrivateKey):
        raise TypeError("Loaded private key is not an Ed25519 key")
    return private_key
def _load_public_key(public_key_b64: str) -> ed25519.Ed25519PublicKey:
    try:
        public_key_bytes = base64.b64decode(public_key_b64, validate=True)
        return ed25519.Ed25519PublicKey.from_public_bytes(public_key_bytes)
    except Exception as e:
        raise ValueError(f"Invalid public key (base64): {e}") from e
def sign_metadata_dict(files_dict: dict[str, Any], private_key_path: str) -> str:
    private_key = _load_private_key(private_key_path)
    data_json = json.dumps(files_dict, sort_keys=True, separators=(",", ":")).encode("utf-8")
    signature = private_key.sign(data_json)
    return base64.b64encode(signature).decode("ascii")
def verify_metadata_signature(
    files_dict: dict[str, Any],
    signature_b64: str,
    public_key_b64: str,
) -> bool:
    public_key = _load_public_key(public_key_b64)
    data_json = json.dumps(files_dict, sort_keys=True, separators=(",", ":")).encode("utf-8")
    signature = base64.b64decode(signature_b64, validate=True)
    try:
        public_key.verify(signature, data_json)
        return True
    except Exception as e:
        raise MetadataSignatureError(f"Metadata signature invalid: {e}") from e
def _iter_project_files(
    root_dir: str,
    *,
    exclude_dirs: set[str],
    exclude_files: set[str],
):
    root_path = Path(root_dir)
    for current_root, dirs, files in os.walk(root_path):
        dirs[:] = [d for d in dirs if d not in exclude_dirs]
        for file in files:
            if file in exclude_files:
                continue
            file_path = Path(current_root) / file
            yield file_path, os.path.relpath(file_path, root_path)
def generate_metadata(
    root_dir: str,
    metadata_file: str = "metadata.json",
    exclude_dirs: list[str] | None = None,
    exclude_files: list[str] | None = None,
    sign: bool = False,
    private_key_path: str | None = None,
    *,
    hash_algorithm: str = "sha256",
) -> dict[str, Any]:
    exclude_dirs_set = set(exclude_dirs or [])
    exclude_files_set = set(exclude_files or [])
    files_metadata: dict[str, dict[str, Any]] = {}
    for file_path, rel_path in _iter_project_files(
        root_dir,
        exclude_dirs=exclude_dirs_set,
        exclude_files=exclude_files_set,
    ):
        file_hash, loc = calculate_hash(str(file_path), algorithm=hash_algorithm)
        files_metadata[rel_path] = {
            "hash": file_hash,
            "lines": loc,
            "algorithm": hash_algorithm,
        }
        logger.debug(f"Indexed {rel_path} ({loc} lines)")
    if sign:
        if not private_key_path:
            raise ValueError("private_key_path must be provided when sign=True")
        signature_b64 = sign_metadata_dict(files_metadata, private_key_path)
        final_data: dict[str, Any] = {
            "signature": signature_b64,
            "files": files_metadata,
        }
    else:
        final_data = files_metadata
    _atomic_write_json(metadata_file, final_data)
    logger.info(f"Successfully generated metadata file at {metadata_file} (signed={sign})")
    return final_data
def verify_integrity(
    root_dir: str,
    metadata_file: str = "metadata.json",
    expect_signed: bool = False,
    public_key_b64: str | None = None,
) -> bool:
    path = Path(metadata_file)
    if not path.exists():
        raise FileNotFoundError(f"{metadata_file} is not a file or directory")
    if not is_valid_json_file(path):
        raise NotJson(f"{metadata_file} is not a valid JSON file")
    with open(path, "r", encoding="utf-8") as f:
        loaded = json.load(f)
    if expect_signed:
        signature_b64 = loaded.get("signature", "")
        files_data = loaded.get("files", {})
        if not signature_b64 or not isinstance(files_data, dict):
            raise IntegrityError("Signed metadata is malformed (missing signature/files)")
        if not public_key_b64:
            raise ValueError("public_key_b64 must be provided when expect_signed=True")
        verify_metadata_signature(files_data, signature_b64, public_key_b64)
        stored_data = files_data
        logger.info("Metadata signature verified: authentic")
    else:
        if not isinstance(loaded, dict):
            raise IntegrityError("Metadata format is invalid")
        stored_data = loaded
    mismatched: list[str] = []
    for file_path, rel_path in _iter_project_files(
        root_dir,
        exclude_dirs=set(),
        exclude_files={os.path.basename(metadata_file)},
    ):
        current_hash, _ = calculate_hash(str(file_path), algorithm=stored_data.get(rel_path, {}).get("algorithm", "sha256"))
        if rel_path not in stored_data:
            logger.warning(f"[NEW] {rel_path}")
            mismatched.append(rel_path)
            continue
        expected_hash = stored_data[rel_path]["hash"]
        if current_hash != expected_hash:
            logger.warning(f"[CHANGED] {rel_path}")
            mismatched.append(rel_path)
        else:
            logger.debug(f"[OK] {rel_path}")
    for rel_path in stored_data.keys():
        if rel_path in {"signature", "files"}:
            continue
        full_path = Path(root_dir) / rel_path
        if not full_path.exists():
            logger.warning(f"[DELETED] {rel_path}")
            mismatched.append(rel_path)
    if not mismatched:
        logger.info("Integrity verified")
        return True
    raise IntegrityError(
        f"Integrity failed: {len(mismatched)} file(s) have been modified, added, or deleted"
    )