from __future__ import annotations

import io
import zipfile
from pathlib import Path
from types import SimpleNamespace

import pytest
from fastapi import HTTPException, UploadFile

from app.gateway.routers import uploads as uploads_router
from deerflow.persistence.upload_policy import MemoryUploadPolicyStore
from deerflow.uploads.policy import classify_upload_filename, default_upload_policy
from deerflow.uploads.validation import UploadContentError, validate_upload_content


def _ooxml_bytes(prefix: str) -> bytes:
    output = io.BytesIO()
    with zipfile.ZipFile(output, "w") as archive:
        archive.writestr("[Content_Types].xml", "<Types />")
        archive.writestr(f"{prefix}content.xml", "<root />")
    return output.getvalue()


@pytest.mark.parametrize(
    ("extension", "content"),
    [
        (".txt", "中文文本".encode()),
        (".md", b"# heading"),
        (".csv", b"a,b\n1,2\n"),
        (".pdf", b"%PDF-1.7\n"),
        (".png", b"\x89PNG\r\n\x1a\nrest"),
        (".jpg", b"\xff\xd8\xff\xe0rest"),
        (".doc", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"),
        (".ppt", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"),
        (".xls", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"),
        (".docx", _ooxml_bytes("word/")),
        (".pptx", _ooxml_bytes("ppt/")),
        (".xlsx", _ooxml_bytes("xl/")),
    ],
)
def test_supported_signatures(extension: str, content: bytes, tmp_path: Path):
    path = tmp_path / f"sample{extension}"
    path.write_bytes(content)
    validate_upload_content(path, extension)


def test_spoofed_extension_is_rejected(tmp_path: Path):
    path = tmp_path / "fake.png"
    path.write_bytes(b"this is not an image")
    with pytest.raises(UploadContentError):
        validate_upload_content(path, ".png")


async def _request_with_policy(policy: dict):
    store = MemoryUploadPolicyStore()
    await store.ensure_seed(policy)
    return SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(upload_policy_store=store)))


def _patch_upload_environment(monkeypatch, uploads_dir: Path, *, sandbox=None):
    uploads_dir.mkdir(parents=True, exist_ok=True)
    monkeypatch.setattr(uploads_router, "ensure_uploads_dir", lambda _thread_id: uploads_dir)
    monkeypatch.setattr(
        uploads_router,
        "get_paths",
        lambda: SimpleNamespace(sandbox_uploads_dir=lambda _thread_id, user_id=None: uploads_dir),
    )

    async def path_user(_thread_id):
        return "user-1"

    monkeypatch.setattr(uploads_router, "aresolve_path_user_id", path_user)
    if sandbox is None:
        provider = SimpleNamespace(uses_thread_data_mounts=True)
    else:
        provider = SimpleNamespace(
            uses_thread_data_mounts=False,
            acquire=lambda _thread_id: "sandbox-1",
            get=lambda _sandbox_id: sandbox,
        )
    monkeypatch.setattr(uploads_router, "get_sandbox_provider", lambda: provider)


@pytest.mark.asyncio
async def test_upload_uses_category_limit_and_returns_category(monkeypatch, tmp_path: Path):
    policy = default_upload_policy()
    policy["groups"]["document"]["max_bytes"] = 3
    request = await _request_with_policy(policy)
    uploads_dir = tmp_path / "uploads"
    _patch_upload_environment(monkeypatch, uploads_dir)
    config = SimpleNamespace(uploads={"auto_convert_documents": False})

    too_large = UploadFile(filename="large.txt", file=io.BytesIO(b"four"))
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, [too_large], config)
    assert exc_info.value.status_code == 413
    assert not (uploads_dir / "large.txt").exists()

    valid = UploadFile(filename="ok.txt", file=io.BytesIO(b"abc"))
    response = await uploads_router.upload_files.__wrapped__("thread-1", request, [valid], config)
    assert response.files[0]["category"] == "document"
    assert response.files[0]["detected_type"] == "txt"
    assert response.policy_revision == 1
    assert (uploads_dir / "ok.txt").read_bytes() == b"abc"


@pytest.mark.asyncio
async def test_batch_validation_preserves_existing_same_name(monkeypatch, tmp_path: Path):
    request = await _request_with_policy(default_upload_policy())
    uploads_dir = tmp_path / "uploads"
    _patch_upload_environment(monkeypatch, uploads_dir)
    (uploads_dir / "same.txt").write_bytes(b"old")
    config = SimpleNamespace(uploads={"auto_convert_documents": False})
    files = [
        UploadFile(filename="same.txt", file=io.BytesIO(b"new")),
        UploadFile(filename="bad.png", file=io.BytesIO(b"not-png")),
    ]
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, files, config)
    assert exc_info.value.status_code == 415
    assert (uploads_dir / "same.txt").read_bytes() == b"old"
    assert not (uploads_dir / "bad.png").exists()


@pytest.mark.asyncio
async def test_remote_sync_failure_rolls_back_committed_files(monkeypatch, tmp_path: Path):
    request = await _request_with_policy(default_upload_policy())
    uploads_dir = tmp_path / "uploads"
    (uploads_dir / "same.txt").parent.mkdir(parents=True)
    (uploads_dir / "same.txt").write_bytes(b"old")

    class FailingSandbox:
        def update_file(self, _path, _content):
            raise OSError("remote write failed")

    _patch_upload_environment(monkeypatch, uploads_dir, sandbox=FailingSandbox())
    config = SimpleNamespace(uploads={"auto_convert_documents": False})
    file = UploadFile(filename="same.txt", file=io.BytesIO(b"new"))
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, [file], config)
    assert exc_info.value.status_code == 500
    assert (uploads_dir / "same.txt").read_bytes() == b"old"


@pytest.mark.asyncio
async def test_count_and_extension_are_enforced_before_writing(monkeypatch, tmp_path: Path):
    policy = default_upload_policy()
    policy["max_files_per_message"] = 2
    request = await _request_with_policy(policy)
    uploads_dir = tmp_path / "uploads"
    _patch_upload_environment(monkeypatch, uploads_dir)
    config = SimpleNamespace(uploads={"auto_convert_documents": False})

    files = [UploadFile(filename=f"{index}.txt", file=io.BytesIO(b"x")) for index in range(3)]
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, files, config)
    assert exc_info.value.status_code == 413

    forbidden = UploadFile(filename="malware.exe", file=io.BytesIO(b"MZ"))
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, [forbidden], config)
    assert exc_info.value.status_code == 415
    assert list(uploads_dir.iterdir()) == []


@pytest.mark.asyncio
async def test_default_accepts_twenty_and_rejects_twenty_one(monkeypatch, tmp_path: Path):
    request = await _request_with_policy(default_upload_policy())
    uploads_dir = tmp_path / "uploads"
    _patch_upload_environment(monkeypatch, uploads_dir)
    config = SimpleNamespace(uploads={"auto_convert_documents": False})

    twenty = [UploadFile(filename=f"ok-{index}.txt", file=io.BytesIO(b"x")) for index in range(20)]
    response = await uploads_router.upload_files.__wrapped__("thread-1", request, twenty, config)
    assert len(response.files) == 20

    twenty_one = [UploadFile(filename=f"no-{index}.txt", file=io.BytesIO(b"x")) for index in range(21)]
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, twenty_one, config)
    assert exc_info.value.status_code == 413
    assert not any((uploads_dir / f"no-{index}.txt").exists() for index in range(21))


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("filename", "prefix"),
    [
        ("document.txt", b"a"),
        ("picture.PNG", b"\x89PNG\r\n\x1a\n"),
        ("table.csv", b"a,b\n"),
    ],
)
async def test_each_category_exact_limit_and_limit_plus_one(filename: str, prefix: bytes, tmp_path: Path):
    policy = default_upload_policy()
    category, limit = classify_upload_filename(filename, policy) or (None, None)
    assert category is not None and limit is not None

    exact_dir = tmp_path / f"exact-{category}"
    exact_dir.mkdir()
    exact_body = prefix + b"x" * (limit - len(prefix))
    _, size, total = await uploads_router._write_upload_file_with_limits(
        UploadFile(filename=filename, file=io.BytesIO(exact_body)),
        uploads_dir=exact_dir,
        display_filename=filename,
        max_single_file_size=limit,
        max_total_size=100 * 1024 * 1024,
        total_size=0,
    )
    assert size == total == limit

    over_dir = tmp_path / f"over-{category}"
    over_dir.mkdir()
    over_body = prefix + b"x" * (limit + 1 - len(prefix))
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router._write_upload_file_with_limits(
            UploadFile(filename=filename, file=io.BytesIO(over_body)),
            uploads_dir=over_dir,
            display_filename=filename,
            max_single_file_size=limit,
            max_total_size=100 * 1024 * 1024,
            total_size=0,
        )
    assert exc_info.value.status_code == 413
    assert not (over_dir / filename).exists()


@pytest.mark.asyncio
async def test_total_size_exact_boundary_and_plus_one(monkeypatch, tmp_path: Path):
    policy = default_upload_policy()
    policy["max_total_bytes"] = 10
    policy["groups"]["document"]["max_bytes"] = 10
    policy["groups"]["image"]["max_bytes"] = 10
    policy["groups"]["spreadsheet"]["max_bytes"] = 5
    policy["large_file"]["threshold_bytes"] = 10
    request = await _request_with_policy(policy)
    uploads_dir = tmp_path / "uploads"
    _patch_upload_environment(monkeypatch, uploads_dir)
    config = SimpleNamespace(uploads={"auto_convert_documents": False})

    exact = [
        UploadFile(filename="first.txt", file=io.BytesIO(b"1234")),
        UploadFile(filename="second.txt", file=io.BytesIO(b"123456")),
    ]
    response = await uploads_router.upload_files.__wrapped__("thread-1", request, exact, config)
    assert len(response.files) == 2

    over = [
        UploadFile(filename="third.txt", file=io.BytesIO(b"1234")),
        UploadFile(filename="fourth.txt", file=io.BytesIO(b"1234567")),
    ]
    with pytest.raises(HTTPException) as exc_info:
        await uploads_router.upload_files.__wrapped__("thread-1", request, over, config)
    assert exc_info.value.status_code == 413
    assert not (uploads_dir / "third.txt").exists()
    assert not (uploads_dir / "fourth.txt").exists()