"""Accept uploads on Core disk and enqueue R2 ingest jobs."""
from __future__ import annotations

import asyncio
import os
import re
from uuid import UUID, uuid4

import math

import structlog
from fastapi import UploadFile
from sqlalchemy.ext.asyncio import AsyncSession

from src.config import Settings, get_settings
from src.core.exceptions import NotFoundError, ValidationError
from src.domain.connectors.repository import ConnectorRepository
from src.domain.projects.models import Project
from src.domain.storage_ingest.models import ArchivePathMode, StorageIngestJob, StorageIngestJobStatus
from src.domain.storage_ingest.paths import build_staging_relative_path
from src.domain.storage_ingest.repository import StorageIngestJobRepository
from src.domain.storage_ingest.schemas import (
    IngestUploadPartResponse,
    IngestUploadSessionResponse,
    StartIngestUploadSessionRequest,
    StorageIngestJobResponse,
    StorageIngestJobStatusResponse,
)
from src.integrations.archive_extract import is_archive_filename
from src.domain.storage_ingest.background import schedule_storage_ingest_fallback
from src.workers.enqueue import enqueue_storage_ingest_job

logger = structlog.get_logger()


def _safe_filename(name: str) -> str:
    base = os.path.basename(name or "") or "upload.bin"
    base = re.sub(r"[^a-zA-Z0-9._-]+", "_", base).strip("._")[:180]
    return base or "upload.bin"


def _normalize_object_key(raw: str | None, filename: str) -> str:
    key = (raw or filename or "upload.bin").strip().replace("\\", "/")
    return key.lstrip("/")


def _poll_url(project_id: UUID, job_id: UUID) -> str:
    return f"/api/projects/{project_id}/storage/ingest-jobs/{job_id}"


def _normalize_archive_path_mode(raw: str | None) -> str:
    mode = (raw or ArchivePathMode.PRESERVE_FOLDERS.value).strip().lower()
    allowed = {ArchivePathMode.PRESERVE_FOLDERS.value, ArchivePathMode.FLATTEN.value}
    if mode not in allowed:
        raise ValidationError("archive_path_mode must be preserve_folders or flatten")
    return mode


class StorageIngestService:
    def __init__(
        self,
        *,
        db: AsyncSession,
        connector_repo: ConnectorRepository,
        job_repo: StorageIngestJobRepository,
        settings: Settings | None = None,
    ) -> None:
        self._db = db
        self._connectors = connector_repo
        self._jobs = job_repo
        self._settings = settings or get_settings()

    async def _validate_r2_connection(self, project: Project, connection_id: UUID):
        conn = await self._connectors.get_connection_for_project(project.id, connection_id)
        if not conn:
            raise NotFoundError("Connector connection not found")
        connector = await self._connectors.get_connector_by_id(conn.connector_id)
        if not connector or connector.key != "cloudflare_r2":
            raise ValidationError("Only cloudflare_r2 connections support storage ingest")
        return conn

    def _staging_root(self) -> str:
        root = os.path.normpath(self._settings.storage_staging_upload_root)
        if not root:
            raise ValidationError("Storage staging is not configured on this instance")
        return root

    def _prepare_upload_meta(
        self,
        *,
        filename: str,
        content_type: str | None,
        size_bytes: int,
        object_key: str | None,
        extract_archive: bool,
    ) -> tuple[str, str, bool, bool, str]:
        if size_bytes <= 0:
            raise ValidationError("Empty file")
        if size_bytes > self._settings.storage_staging_max_bytes:
            raise ValidationError(
                f"File exceeds staging limit ({self._settings.storage_staging_max_bytes} bytes)"
            )
        safe_name = _safe_filename(filename)
        ct = content_type or "application/octet-stream"
        is_archive = is_archive_filename(safe_name, ct)
        should_extract = bool(extract_archive and is_archive)
        normalized_key = _normalize_object_key(object_key, safe_name)
        return safe_name, ct, is_archive, should_extract, normalized_key

    async def start_upload_session(
        self,
        project: Project,
        body: StartIngestUploadSessionRequest,
    ) -> IngestUploadSessionResponse:
        await self._validate_r2_connection(project, body.connection_id)
        root = self._staging_root()
        safe_name, _, is_archive, should_extract, normalized_key = self._prepare_upload_meta(
            filename=body.filename,
            content_type=None,
            size_bytes=body.size_bytes,
            object_key=body.object_key,
            extract_archive=body.extract_archive,
        )
        path_mode = _normalize_archive_path_mode(body.archive_path_mode)
        chunk_size = self._settings.storage_ingest_chunk_bytes
        if chunk_size <= 0:
            raise ValidationError("storage_ingest_chunk_bytes must be positive")

        job_id = uuid4()
        rel = build_staging_relative_path(
            project.workspace_id,
            project.id,
            safe_name,
            ingest_id=job_id,
        )
        abs_path = os.path.join(root, rel.replace("/", os.sep))
        os.makedirs(os.path.dirname(abs_path), exist_ok=True)

        def _touch() -> None:
            with open(abs_path, "wb"):
                pass

        await asyncio.to_thread(_touch)

        row = StorageIngestJob(
            id=job_id,
            project_id=project.id,
            connection_id=body.connection_id,
            bucket_name=body.bucket_name.strip(),
            status=StorageIngestJobStatus.UPLOADING.value,
            original_filename=safe_name,
            staging_relative_path=rel.replace("\\", "/"),
            object_key=normalized_key,
            extract_archive=should_extract,
            is_archive=is_archive,
            archive_path_mode=path_mode,
            expected_size_bytes=body.size_bytes,
            received_bytes=0,
            next_part_number=1,
        )
        await self._jobs.create(row)
        await self._db.commit()

        total_parts = max(1, math.ceil(body.size_bytes / chunk_size))
        logger.info(
            "storage_ingest_session_started",
            job_id=str(job_id),
            project_id=str(project.id),
            bytes=body.size_bytes,
            parts=total_parts,
        )
        return IngestUploadSessionResponse(
            job_id=str(job_id),
            status=row.status,
            poll_url=_poll_url(project.id, job_id),
            bucket_name=row.bucket_name,
            extract_archive=should_extract,
            is_archive=is_archive,
            chunk_size_bytes=chunk_size,
            total_parts=total_parts,
            expected_size_bytes=body.size_bytes,
        )

    async def append_upload_part(
        self,
        project_id: UUID,
        job_id: UUID,
        part_number: int,
        chunk: bytes,
    ) -> IngestUploadPartResponse:
        if part_number < 1:
            raise ValidationError("part_number must be >= 1")
        if not chunk:
            raise ValidationError("Empty upload part")

        row = await self._jobs.get_for_project(job_id, project_id)
        if not row:
            raise NotFoundError("Ingest job not found")
        if row.status != StorageIngestJobStatus.UPLOADING.value:
            raise ValidationError("upload_session_not_accepting_parts")

        chunk_size = self._settings.storage_ingest_chunk_bytes
        expected = row.expected_size_bytes or 0
        if part_number != row.next_part_number:
            raise ValidationError(
                f"expected_part_{row.next_part_number}_got_{part_number}"
            )
        if row.received_bytes + len(chunk) > expected:
            raise ValidationError("part_exceeds_declared_file_size")

        is_last_part = part_number == max(1, math.ceil(expected / chunk_size))
        if not is_last_part and len(chunk) != chunk_size:
            raise ValidationError(f"part_must_be_{chunk_size}_bytes")
        if is_last_part and len(chunk) > chunk_size:
            raise ValidationError("last_part_too_large")

        root = self._staging_root()
        abs_path = os.path.join(root, row.staging_relative_path.replace("/", os.sep))

        def _append() -> None:
            with open(abs_path, "ab") as out:
                out.write(chunk)

        await asyncio.to_thread(_append)

        row.received_bytes += len(chunk)
        row.next_part_number += 1
        await self._jobs.save(row)
        await self._db.commit()

        complete = row.received_bytes >= expected
        return IngestUploadPartResponse(
            job_id=str(row.id),
            part_number=part_number,
            received_bytes=row.received_bytes,
            next_part_number=row.next_part_number,
            complete=complete,
        )

    async def complete_upload_session(
        self,
        project_id: UUID,
        job_id: UUID,
    ) -> StorageIngestJobResponse:
        row = await self._jobs.get_for_project(job_id, project_id)
        if not row:
            raise NotFoundError("Ingest job not found")
        if row.status != StorageIngestJobStatus.UPLOADING.value:
            raise ValidationError("upload_session_not_ready_to_complete")

        expected = row.expected_size_bytes or 0
        if row.received_bytes != expected:
            raise ValidationError(
                f"incomplete_upload: received {row.received_bytes} of {expected} bytes"
            )

        root = self._staging_root()
        abs_path = os.path.join(root, row.staging_relative_path.replace("/", os.sep))
        if not os.path.isfile(abs_path):
            raise ValidationError("staging_file_missing")

        row.status = StorageIngestJobStatus.PENDING.value
        await self._jobs.save(row)
        await self._db.commit()

        enqueued = await enqueue_storage_ingest_job(str(job_id))
        if not enqueued:
            row.status = StorageIngestJobStatus.FAILED.value
            row.error = "worker_enqueue_failed"
            await self._jobs.save(row)
            await self._db.commit()
            raise ValidationError("Could not enqueue ingest job — is the worker running?")

        # The API process that wrote the file must start R2 upload immediately.
        # arq often runs on a host that cannot see local staging and used to
        # mark the job failed before this fallback ran.
        schedule_storage_ingest_fallback(str(job_id), delay_seconds=0)

        logger.info(
            "storage_ingest_session_completed",
            job_id=str(job_id),
            project_id=str(project_id),
            bytes=row.received_bytes,
        )
        return StorageIngestJobResponse(
            job_id=str(row.id),
            status=row.status,
            poll_url=_poll_url(project_id, job_id),
            bucket_name=row.bucket_name,
            extract_archive=row.extract_archive,
            is_archive=row.is_archive,
        )

    async def accept_upload(
        self,
        project: Project,
        *,
        connection_id: UUID,
        bucket_name: str,
        file: UploadFile,
        object_key: str | None,
        extract_archive: bool,
        archive_path_mode: str = ArchivePathMode.PRESERVE_FOLDERS.value,
    ) -> StorageIngestJobResponse:
        bucket = bucket_name.strip()
        if not bucket:
            raise ValidationError("bucket_name is required")

        await self._validate_r2_connection(project, connection_id)
        root = self._staging_root()

        filename = _safe_filename(file.filename or "upload.bin")
        content_type = file.content_type or "application/octet-stream"
        content = await file.read()
        safe_name, _, is_archive, should_extract, normalized_key = self._prepare_upload_meta(
            filename=filename,
            content_type=content_type,
            size_bytes=len(content),
            object_key=object_key,
            extract_archive=extract_archive,
        )
        path_mode = _normalize_archive_path_mode(archive_path_mode)
        if len(content) > self._settings.storage_ingest_single_post_max_bytes:
            raise ValidationError(
                "File is too large for a single upload request. "
                "Use the chunked ingest-upload session API instead."
            )

        job_id = uuid4()
        rel = build_staging_relative_path(
            project.workspace_id,
            project.id,
            safe_name,
            ingest_id=job_id,
        )
        abs_path = os.path.join(root, rel.replace("/", os.sep))
        os.makedirs(os.path.dirname(abs_path), exist_ok=True)

        def _write() -> None:
            with open(abs_path, "wb") as out:
                out.write(content)

        await asyncio.to_thread(_write)

        row = StorageIngestJob(
            id=job_id,
            project_id=project.id,
            connection_id=connection_id,
            bucket_name=bucket,
            status=StorageIngestJobStatus.PENDING.value,
            original_filename=safe_name,
            staging_relative_path=rel.replace("\\", "/"),
            object_key=normalized_key,
            extract_archive=should_extract,
            is_archive=is_archive,
            archive_path_mode=path_mode,
        )
        await self._jobs.create(row)
        await self._db.commit()

        enqueued = await enqueue_storage_ingest_job(str(job_id))
        if not enqueued:
            row.status = StorageIngestJobStatus.FAILED.value
            row.error = "worker_enqueue_failed"
            await self._jobs.save(row)
            await self._db.commit()
            raise ValidationError("Could not enqueue ingest job — is the worker running?")

        schedule_storage_ingest_fallback(str(job_id), delay_seconds=0)

        logger.info(
            "storage_ingest_accepted",
            job_id=str(job_id),
            project_id=str(project.id),
            bucket=bucket,
            bytes=len(content),
            extract=should_extract,
        )

        return StorageIngestJobResponse(
            job_id=str(job_id),
            status=row.status,
            poll_url=_poll_url(project.id, job_id),
            bucket_name=bucket,
            extract_archive=should_extract,
            is_archive=is_archive,
        )

    async def get_job_status(
        self,
        project_id: UUID,
        job_id: UUID,
    ) -> StorageIngestJobStatusResponse:
        row = await self._jobs.get_for_project(job_id, project_id)
        if not row:
            raise NotFoundError("Ingest job not found")
        return StorageIngestJobStatusResponse(
            job_id=str(row.id),
            status=row.status,
            bucket_name=row.bucket_name,
            original_filename=row.original_filename,
            extract_archive=row.extract_archive,
            is_archive=row.is_archive,
            archive_path_mode=row.archive_path_mode,
            uploaded_count=row.uploaded_count,
            total_entries=row.total_entries,
            skipped_count=row.skipped_count,
            expected_size_bytes=row.expected_size_bytes,
            received_bytes=row.received_bytes,
            error=row.error,
            result=row.result_json,
        )

    async def retry_ingest_job(
        self,
        project_id: UUID,
        job_id: UUID,
    ) -> StorageIngestJobResponse:
        """Re-enqueue a stuck pending/failed job when the staging file is still on disk."""
        row = await self._jobs.get_for_project(job_id, project_id)
        if not row:
            raise NotFoundError("Ingest job not found")
        if row.status not in (
            StorageIngestJobStatus.PENDING.value,
            StorageIngestJobStatus.FAILED.value,
        ):
            raise ValidationError(f"job_not_retryable: status={row.status}")

        expected = row.expected_size_bytes or 0
        if expected and row.received_bytes != expected:
            raise ValidationError("upload_incomplete")

        root = self._staging_root()
        abs_path = os.path.join(root, row.staging_relative_path.replace("/", os.sep))
        if not os.path.isfile(abs_path):
            raise ValidationError("staging_file_missing")

        if row.status == StorageIngestJobStatus.FAILED.value:
            row.status = StorageIngestJobStatus.PENDING.value
            row.error = None
            await self._jobs.save(row)
            await self._db.commit()

        enqueued = await enqueue_storage_ingest_job(str(job_id), retry=True)
        if not enqueued:
            raise ValidationError("worker_enqueue_failed")

        schedule_storage_ingest_fallback(str(job_id), delay_seconds=0)

        logger.info("storage_ingest_job_requeued", job_id=str(job_id), project_id=str(project_id))
        return StorageIngestJobResponse(
            job_id=str(row.id),
            status=row.status,
            poll_url=_poll_url(project_id, job_id),
            bucket_name=row.bucket_name,
            extract_archive=row.extract_archive,
            is_archive=row.is_archive,
        )