from __future__ import annotations

import fcntl
import json
import os
import sqlite3
from collections.abc import Iterator
from pathlib import Path
from typing import Any

from .crypto import (
    ENCRYPTION_VERSION,
    EncryptionError,
    decrypt_file_bytes,
    decrypt_json,
    encrypt_file,
    encrypt_json,
    file_aad,
    is_encrypted_file,
    iter_decrypted_file,
    metadata_aad,
)


class EncryptedArtifactStore:
    def __init__(self, storage_dir: Path, master_key: bytes | None, key_version: int = 1):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.master_key = master_key
        self.key_version = int(key_version)

    @property
    def configured(self) -> bool:
        return self.master_key is not None

    def require_key(self) -> bytes:
        if self.master_key is None:
            raise EncryptionError(
                "KIMU_ENCRYPTION_KEY_REQUIRED",
                "Encrypted artifact storage requires KIMU_MASTER_KEY.",
            )
        return self.master_key

    def body_path(self, row: Any) -> Path:
        return self.storage_dir / row["stored_name"]

    def encode_private_metadata(
        self,
        file_id: str,
        sha256: str,
        text_content: str,
        zip_inventory_json: str,
        warnings_json: str,
        key_version: int | None = None,
    ) -> bytes:
        version = int(key_version or self.key_version)
        value = {
            "text_content": text_content or "",
            "zip_inventory": json.loads(zip_inventory_json or "[]"),
            "warnings": json.loads(warnings_json or "[]"),
        }
        return encrypt_json(
            value,
            self.require_key(),
            metadata_aad(file_id, sha256, version),
        )

    def private_metadata(self, row: Any) -> dict[str, Any]:
        encryption_version = int(row["encryption_version"] or 0)
        if encryption_version <= 0:
            return {
                "text_content": row["text_content"] or "",
                "zip_inventory": json.loads(row["zip_inventory_json"] or "[]"),
                "warnings": json.loads(row["warnings_json"] or "[]"),
            }
        version = int(row["key_version"] or 1)
        value = decrypt_json(
            row["private_metadata_cipher"],
            self.require_key(),
            metadata_aad(row["id"], row["sha256"], version),
        )
        value.setdefault("text_content", "")
        value.setdefault("zip_inventory", [])
        value.setdefault("warnings", [])
        return value

    def encrypt_staged(
        self, staged: Path, final_path: Path, file_id: str, sha256: str
    ) -> int:
        return encrypt_file(
            staged,
            final_path,
            self.require_key(),
            file_aad(file_id, sha256, self.key_version),
        )

    def iter_body(self, row: Any) -> Iterator[bytes]:
        path = self.body_path(row)
        if int(row["encryption_version"] or 0) <= 0:
            with path.open("rb") as handle:
                for block in iter(lambda: handle.read(1024 * 1024), b""):
                    yield block
            return
        yield from iter_decrypted_file(
            path,
            self.require_key(),
            file_aad(row["id"], row["sha256"], int(row["key_version"] or 1)),
        )

    def body_bytes(self, row: Any, maximum_bytes: int | None = None) -> bytes:
        path = self.body_path(row)
        if int(row["encryption_version"] or 0) <= 0:
            data = path.read_bytes()
            if maximum_bytes is not None and len(data) > maximum_bytes:
                raise EncryptionError(
                    "KIMU_DECRYPT_LIMIT",
                    "The file body exceeds the permitted in-memory limit.",
                )
            return data
        return decrypt_file_bytes(
            path,
            self.require_key(),
            file_aad(row["id"], row["sha256"], int(row["key_version"] or 1)),
            maximum_bytes,
        )

    def migrate_legacy_files(self, db: sqlite3.Connection) -> dict[str, int]:
        if not self.configured:
            return {"migrated": 0, "missing": 0}
        lock_path = self.storage_dir / ".encryption-migration.lock"
        with lock_path.open("a+b") as lock_handle:
            fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX)
            try:
                return self._migrate_legacy_files_unlocked(db)
            finally:
                fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN)

    def _migrate_legacy_files_unlocked(self, db: sqlite3.Connection) -> dict[str, int]:
        migrated = 0
        missing = 0
        rows = db.execute(
            "SELECT * FROM files WHERE encryption_version=0 AND deleted_at IS NULL"
        ).fetchall()
        for row in rows:
            path = self.body_path(row)
            if not path.is_file():
                missing += 1
                continue
            if is_encrypted_file(path):
                stored_size = path.stat().st_size
            else:
                temporary = path.with_name(f".{path.name}.legacy")
                os.replace(path, temporary)
                try:
                    stored_size = self.encrypt_staged(
                        temporary, path, row["id"], row["sha256"]
                    )
                except Exception:
                    if not path.exists() and temporary.exists():
                        os.replace(temporary, path)
                    raise
                finally:
                    temporary.unlink(missing_ok=True)
            cipher = self.encode_private_metadata(
                row["id"],
                row["sha256"],
                row["text_content"] or "",
                row["zip_inventory_json"] or "[]",
                row["warnings_json"] or "[]",
            )
            db.execute(
                """
                UPDATE files SET private_metadata_cipher=?,text_content='',
                  zip_inventory_json='[]',warnings_json='[]',encryption_version=?,
                  key_version=?,stored_size_bytes=? WHERE id=?
                """,
                (
                    cipher,
                    ENCRYPTION_VERSION,
                    self.key_version,
                    stored_size,
                    row["id"],
                ),
            )
            migrated += 1
        return {"migrated": migrated, "missing": missing}