from __future__ import annotations

import base64
import hashlib
import json
import mimetypes
import os
import re
import unicodedata
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Any


TEXT_EXTENSIONS = {
    ".txt", ".md", ".markdown", ".rst", ".csv", ".tsv", ".json", ".jsonl",
    ".yaml", ".yml", ".toml", ".ini", ".cfg", ".conf", ".log", ".xml", ".html",
    ".htm", ".css", ".scss", ".less", ".js", ".jsx", ".ts", ".tsx", ".mjs",
    ".cjs", ".py", ".pyi", ".java", ".kt", ".kts", ".c", ".h", ".cc", ".cpp",
    ".hpp", ".cs", ".go", ".rs", ".php", ".rb", ".swift", ".sh", ".bash",
    ".zsh", ".fish", ".ps1", ".sql", ".graphql", ".gql", ".vue", ".svelte",
    ".dockerfile", ".gitignore", ".env", ".properties", ".gradle", ".cmake",
    ".lua", ".r", ".tex", ".bib", ".ipynb",
}
IMAGE_MIME_PREFIX = "image/"
ARCHIVE_EXTENSIONS = {".zip"}
MAX_ZIP_MEMBERS = 2000
MAX_ZIP_UNCOMPRESSED = 200 * 1024 * 1024
MAX_ZIP_MEMBER = 32 * 1024 * 1024
MAX_ZIP_RATIO = 120
MAX_TEXT_MEMBER = 256 * 1024
MAX_TEXT_TOTAL = 2 * 1024 * 1024
MAX_IMAGE_INLINE = 12 * 1024 * 1024


class FileValidationError(ValueError):
    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code


@dataclass(frozen=True)
class FileInspection:
    original_name: str
    safe_name: str
    mime_type: str
    size_bytes: int
    sha256: str
    kind: str
    text_content: str
    zip_inventory: list[dict[str, Any]]
    warnings: list[str]

    def database_values(self) -> dict[str, Any]:
        return {
            "original_name": self.original_name,
            "safe_name": self.safe_name,
            "mime_type": self.mime_type,
            "size_bytes": self.size_bytes,
            "sha256": self.sha256,
            "kind": self.kind,
            "text_content": self.text_content,
            "zip_inventory_json": json.dumps(self.zip_inventory, ensure_ascii=False, separators=(",", ":")),
            "warnings_json": json.dumps(self.warnings, ensure_ascii=False, separators=(",", ":")),
        }


def safe_filename(name: str) -> str:
    leaf = Path(str(name or "upload")).name
    leaf = unicodedata.normalize("NFKC", leaf)
    leaf = re.sub(r"[\x00-\x1f\x7f]+", "", leaf)
    leaf = re.sub(r"[^A-Za-z0-9._()\- ]+", "_", leaf).strip(" .")
    if not leaf:
        leaf = "upload"
    if len(leaf) > 180:
        suffix = Path(leaf).suffix[:20]
        leaf = leaf[: 180 - len(suffix)] + suffix
    return leaf


def sniff_mime(path: Path, supplied: str | None, original_name: str) -> str:
    head = path.read_bytes()[:16]
    if head.startswith(b"PK\x03\x04"):
        return "application/zip"
    if head.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if head[:3] == b"\xff\xd8\xff":
        return "image/jpeg"
    if head.startswith(b"GIF87a") or head.startswith(b"GIF89a"):
        return "image/gif"
    if head.startswith(b"RIFF") and head[8:12] == b"WEBP":
        return "image/webp"
    guessed = mimetypes.guess_type(original_name)[0]
    supplied = (supplied or "").split(";", 1)[0].strip().lower()
    if supplied and supplied != "application/octet-stream":
        return supplied
    return guessed or "application/octet-stream"


def is_text_name(name: str) -> bool:
    lower = name.lower()
    if Path(lower).suffix in TEXT_EXTENSIONS:
        return True
    return Path(lower).name in {"dockerfile", "makefile", "license", "readme", "procfile"}


def _decode_text(data: bytes) -> str:
    if b"\x00" in data[:4096]:
        return ""
    for encoding in ("utf-8", "utf-8-sig", "utf-16", "latin-1"):
        try:
            return data.decode(encoding)
        except UnicodeDecodeError:
            continue
    return ""


def _unsafe_zip_name(name: str) -> bool:
    normalised = name.replace("\\", "/")
    if not normalised or normalised.startswith("/") or re.match(r"^[A-Za-z]:/", normalised):
        return True
    parts = [part for part in normalised.split("/") if part not in {"", "."}]
    return any(part == ".." for part in parts)


def inspect_zip(path: Path) -> tuple[list[dict[str, Any]], str, list[str]]:
    inventory: list[dict[str, Any]] = []
    extracted_sections: list[str] = []
    warnings: list[str] = []
    total_uncompressed = 0
    total_text = 0
    try:
        archive = zipfile.ZipFile(path)
    except (zipfile.BadZipFile, OSError) as exc:
        raise FileValidationError("KIMU_ZIP_INVALID", "The ZIP archive is invalid or unreadable.") from exc
    with archive:
        infos = archive.infolist()
        if len(infos) > MAX_ZIP_MEMBERS:
            raise FileValidationError("KIMU_ZIP_TOO_MANY_MEMBERS", f"The ZIP contains more than {MAX_ZIP_MEMBERS} members.")
        for info in infos:
            name = info.filename.replace("\\", "/")
            if _unsafe_zip_name(name):
                raise FileValidationError("KIMU_ZIP_UNSAFE_PATH", f"Unsafe ZIP member path: {name[:160]}")
            unix_mode = (info.external_attr >> 16) & 0xFFFF
            is_symlink = (unix_mode & 0o170000) == 0o120000
            if is_symlink:
                raise FileValidationError("KIMU_ZIP_SYMLINK_BLOCKED", f"ZIP symbolic links are not permitted: {name[:160]}")
            total_uncompressed += int(info.file_size)
            if info.file_size > MAX_ZIP_MEMBER:
                raise FileValidationError("KIMU_ZIP_MEMBER_TOO_LARGE", f"ZIP member is too large: {name[:160]}")
            if total_uncompressed > MAX_ZIP_UNCOMPRESSED:
                raise FileValidationError("KIMU_ZIP_EXPANSION_LIMIT", "The ZIP exceeds the safe uncompressed-size limit.")
            ratio = (info.file_size / max(info.compress_size, 1)) if info.file_size else 0
            if ratio > MAX_ZIP_RATIO and info.file_size > 1024 * 1024:
                raise FileValidationError("KIMU_ZIP_RATIO_LIMIT", f"Suspicious ZIP compression ratio: {name[:160]}")
            item = {
                "name": name,
                "size": int(info.file_size),
                "compressed_size": int(info.compress_size),
                "directory": info.is_dir(),
                "text": bool(not info.is_dir() and is_text_name(name)),
            }
            inventory.append(item)
            if info.is_dir() or not item["text"] or total_text >= MAX_TEXT_TOTAL:
                continue
            read_limit = min(int(info.file_size), MAX_TEXT_MEMBER)
            try:
                with archive.open(info, "r") as member:
                    data = member.read(read_limit + 1)
            except (RuntimeError, OSError, zipfile.BadZipFile):
                warnings.append(f"Could not read {name}")
                continue
            truncated = len(data) > read_limit
            text = _decode_text(data[:read_limit])
            if not text:
                continue
            remaining = MAX_TEXT_TOTAL - total_text
            text = text[:remaining]
            total_text += len(text.encode("utf-8", errors="ignore"))
            extracted_sections.append(f"\n--- ZIP member: {name} ---\n{text}")
            if truncated:
                warnings.append(f"Text preview truncated: {name}")
    return inventory, "".join(extracted_sections).strip(), warnings


def inspect_stored_file(path: Path, original_name: str, supplied_mime: str | None = None) -> FileInspection:
    size = path.stat().st_size
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    mime = sniff_mime(path, supplied_mime, original_name)
    safe = safe_filename(original_name)
    suffix = Path(safe).suffix.lower()
    inventory: list[dict[str, Any]] = []
    text = ""
    warnings: list[str] = []
    if mime == "application/zip" or suffix in ARCHIVE_EXTENSIONS:
        kind = "zip"
        inventory, text, warnings = inspect_zip(path)
    elif mime.startswith(IMAGE_MIME_PREFIX):
        kind = "image"
    elif is_text_name(safe) or mime.startswith("text/") or mime in {
        "application/json", "application/xml", "application/javascript",
        "application/x-javascript", "application/sql",
    }:
        kind = "text"
        data = path.read_bytes()[:MAX_TEXT_TOTAL + 1]
        if len(data) > MAX_TEXT_TOTAL:
            warnings.append("Text preview truncated.")
        text = _decode_text(data[:MAX_TEXT_TOTAL])
        if not text:
            kind = "binary"
    else:
        kind = "binary"
    return FileInspection(
        original_name=original_name,
        safe_name=safe,
        mime_type=mime,
        size_bytes=size,
        sha256=digest.hexdigest(),
        kind=kind,
        text_content=text,
        zip_inventory=inventory,
        warnings=warnings,
    )


def attachment_content(
    file_row: Any,
    stored_path: Path | None = None,
    *,
    body_bytes: bytes | None = None,
    private_metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
    kind = file_row["kind"]
    def row_value(key: str, default: Any) -> Any:
        try:
            return file_row[key]
        except (KeyError, IndexError):
            return default
    metadata = private_metadata or {
        "text_content": row_value("text_content", "") or "",
        "zip_inventory": json.loads(row_value("zip_inventory_json", "[]") or "[]"),
        "warnings": json.loads(row_value("warnings_json", "[]") or "[]"),
    }
    if kind == "image" and int(file_row["size_bytes"]) <= MAX_IMAGE_INLINE:
        if body_bytes is None and stored_path is not None:
            body_bytes = stored_path.read_bytes()
        if body_bytes is not None:
            encoded = base64.b64encode(body_bytes).decode("ascii")
            return {
                "type": "image_url",
                "image_url": {"url": f"data:{file_row['mime_type']};base64,{encoded}"},
            }
    text = str(metadata.get("text_content") or "")
    inventory = metadata.get("zip_inventory") or []
    if kind == "zip":
        inventory_lines = "\n".join(
            f"- {item['name']} ({item['size']} bytes)" for item in inventory[:500]
        )
        body = (
            f"Attached ZIP: {file_row['original_name']}\n"
            f"SHA-256: {file_row['sha256']}\n"
            f"Safe inventory ({len(inventory)} members):\n{inventory_lines}"
        )
        if text:
            body += f"\n\nExtracted textual content:\n{text}"
    elif text:
        body = (
            f"Attached file: {file_row['original_name']}\n"
            f"MIME: {file_row['mime_type']}\nSHA-256: {file_row['sha256']}\n\n{text}"
        )
    else:
        body = (
            f"Attached file metadata only: {file_row['original_name']}\n"
            f"MIME: {file_row['mime_type']}\nSize: {file_row['size_bytes']} bytes\n"
            f"SHA-256: {file_row['sha256']}\n"
            "Binary content was retained privately but was not inserted into the model context."
        )
    return {"type": "text", "text": body}