import os
import shutil
import pytest
import hashlib
import tempfile
from pathlib import Path
from httpx import AsyncClient, ASGITransport

from app.main import app
from app.core.database import init_db, AsyncSessionLocal
from app.core.config import settings
from app.storage.local import storage_provider
from app.models.models import Client, ClientCredential, BackupSession, BackupFile
from app.core.security import hash_token

@pytest.mark.asyncio
async def test_chunk_upload_interruption_and_resumption():
    """
    CRITICAL PROOF-OF-CONCEPT TEST (#34):
    1. Create a large test binary file with random bytes and known SHA-256.
    2. Start upload session.
    3. Upload first 50% of chunks.
    4. Simulate client disconnect / network drop.
    5. Reconnect client, query session status, ensure server reports already received chunks.
    6. Upload ONLY the remaining 50% of chunks.
    7. Complete upload session.
    8. Verify server-assembled file SHA-256 matches exact original hash.
    """
    await init_db()

    transport = ASGITransport(app=app)
    async with AsyncClient(transport=transport, base_url="http://test") as ac:
        # 1. Setup client and credentials in database
        async with AsyncSessionLocal() as db:
            client = Client(
                client_code="CLIENT-TEST-01",
                name="Windows 11 Test Machine",
                hostname="WIN11-PROD",
                os_info="Windows 11 Pro",
                agent_version="1.0.0",
                status="ONLINE"
            )
            db.add(client)
            await db.flush()

            device_id = "dev-test-uuid-001"
            device_token = "secret-token-test-12345"
            token_hash = hash_token(device_token)

            cred = ClientCredential(
                client_id=client.id,
                device_id=device_id,
                token_hash=token_hash,
                name="Test Agent Credential"
            )
            db.add(cred)
            await db.commit()

        headers = {
            "X-Device-Id": device_id,
            "X-Device-Token": device_token
        }

        # 2. Generate a 12 MB test file (with 4 MB chunks -> 3 chunks total)
        chunk_size = 4 * 1024 * 1024
        file_size = 12 * 1024 * 1024
        test_bytes = os.urandom(file_size)
        expected_sha256 = hashlib.sha256(test_bytes).hexdigest()

        # 3. Initialize upload session
        init_resp = await ac.post("/api/upload/session", json={
            "filename": "database_production.bak",
            "file_size": file_size,
            "sha256": expected_sha256,
            "chunk_size": chunk_size
        }, headers=headers)
        assert init_resp.status_code == 200, f"Init failed: {init_resp.text}"
        session_data = init_resp.json()
        session_code = session_data["session_code"]
        assert session_data["total_chunks"] == 3
        assert session_data["received_chunks"] == []

        # 4. Upload Chunk 0 (0MB to 4MB)
        chunk_0_data = test_bytes[0 : chunk_size]
        chunk_0_hash = hashlib.sha256(chunk_0_data).hexdigest()
        c0_resp = await ac.post(
            f"/api/upload/{session_code}/chunk?chunk_index=0&chunk_sha256={chunk_0_hash}",
            content=chunk_0_data,
            headers={**headers, "Content-Type": "application/octet-stream"}
        )
        assert c0_resp.status_code == 200
        assert c0_resp.json()["is_received"] == True

        # 5. SIMULATE NETWORK CUT / INTERRUPTION!
        # Client disconnects. We now re-query status as if reconnecting later.
        status_resp = await ac.get(f"/api/upload/{session_code}/status", headers=headers)
        assert status_resp.status_code == 200
        status_data = status_resp.json()
        assert status_data["received_chunks"] == [0]
        assert status_data["missing_chunks"] == [1, 2]

        # Or re-call /session endpoint: it should return already received chunks!
        resume_resp = await ac.post("/api/upload/session", json={
            "filename": "database_production.bak",
            "file_size": file_size,
            "sha256": expected_sha256,
            "chunk_size": chunk_size
        }, headers=headers)
        assert resume_resp.status_code == 200
        assert resume_resp.json()["received_chunks"] == [0]

        # 6. Upload Chunk 1 (4MB to 8MB)
        chunk_1_data = test_bytes[chunk_size : 2 * chunk_size]
        chunk_1_hash = hashlib.sha256(chunk_1_data).hexdigest()
        c1_resp = await ac.post(
            f"/api/upload/{session_code}/chunk?chunk_index=1&chunk_sha256={chunk_1_hash}",
            content=chunk_1_data,
            headers={**headers, "Content-Type": "application/octet-stream"}
        )
        assert c1_resp.status_code == 200

        # 7. Upload Chunk 2 (8MB to 12MB)
        chunk_2_data = test_bytes[2 * chunk_size : 3 * chunk_size]
        chunk_2_hash = hashlib.sha256(chunk_2_data).hexdigest()
        c2_resp = await ac.post(
            f"/api/upload/{session_code}/chunk?chunk_index=2&chunk_sha256={chunk_2_hash}",
            content=chunk_2_data,
            headers={**headers, "Content-Type": "application/octet-stream"}
        )
        assert c2_resp.status_code == 200

        # 8. Complete session and trigger assembly & SHA-256 integrity verification
        complete_resp = await ac.post(f"/api/upload/{session_code}/complete", headers=headers)
        assert complete_resp.status_code == 200, f"Complete failed: {complete_resp.text}"
        complete_data = complete_resp.json()
        assert complete_data["status"] == "SUCCESS"
        assert complete_data["sha256"] == expected_sha256
        assert complete_data["file_size"] == file_size

        # 9. Verify physical file on disk and its SHA-256 directly
        assembled_file_path = await storage_provider.get_file_path(complete_data["relative_path"])
        assert os.path.exists(assembled_file_path)
        with open(assembled_file_path, "rb") as f:
            disk_bytes = f.read()
        assert len(disk_bytes) == file_size
        assert hashlib.sha256(disk_bytes).hexdigest() == expected_sha256
        print("\n>>> PoC Test Passed: Interrupted chunk upload successfully resumed and verified bit-for-bit!")