use std::{convert::TryInto, io::Cursor, panic, sync::Arc, time::Duration};

use async_trait::async_trait;
use chrono::Utc;
use diesel::{pg::upsert::excluded, prelude::*};
use image::{GenericImageView, ImageFormat, ImageReader};
use pdfium_render::prelude::*;
use serde::Deserialize;
use serde_json::{json, Map, Value};
use tokio::task;
use tracing::{error, info, warn};
use uuid::Uuid;

use crate::{
    documents::asset::delete_asset,
    jobs::JOB_GENERATE_THUMBNAILS,
    models::{
        Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
        NewDocumentAssetObject,
    },
    schema::{document_asset_objects, document_assets, document_versions, documents},
    state::AppState,
    storage::TenantStorage,
    utils::storage_paths::document_asset_object_key,
};

use super::{
    analyze::determine_thumbnail_support,
    fetch_version_object,
    handle_fetch_error,
    JobExecution,
    JobHandler,
};

const THUMBNAIL_WIDTH: u32 = 512;
const THUMBNAIL_HEIGHT: u32 = 512;
const PREVIEW_WIDTH: u32 = THUMBNAIL_WIDTH * 4;
const PREVIEW_HEIGHT: u32 = THUMBNAIL_HEIGHT * 4;
const THUMBNAIL_ASSET_TYPE: &str = "thumbnail";
const PREVIEW_ASSET_TYPE: &str = "preview";

#[derive(Debug, Deserialize)]
struct ThumbnailPayload {
    document_id: Uuid,
    document_version_id: Uuid,
    #[serde(default)]
    force: bool,
}

pub struct GenerateThumbnailsJob;

impl GenerateThumbnailsJob {
    pub fn new() -> Self {
        Self
    }
}

#[async_trait]
impl JobHandler for GenerateThumbnailsJob {
    fn job_type(&self) -> &'static str {
        JOB_GENERATE_THUMBNAILS
    }

    async fn handle(
        &self,
        state: Arc<AppState>,
        job: crate::models::Job,
        storage: TenantStorage,
    ) -> JobExecution {
        let payload: ThumbnailPayload = match serde_json::from_value(job.payload.clone()) {
            Ok(p) => p,
            Err(err) => {
                return JobExecution::Failed {
                    error: format!("invalid thumbnail payload: {err}"),
                }
            }
        };

        let state_clone = state.clone();
        let initial =
            match task::spawn_blocking(move || load_thumbnail_context(state_clone, &payload)).await
            {
                Ok(Ok(ctx)) => ctx,
                Ok(Err(err)) => {
                    warn!(job_id = %job.id, error = %err, "thumbnail job will retry");
                    return JobExecution::Retry {
                        delay: Duration::from_secs(30),
                        error: err,
                    };
                }
                Err(join_err) => {
                    error!(job_id = %job.id, error = %join_err, "thumbnail task panicked");
                    return JobExecution::Retry {
                        delay: Duration::from_secs(60),
                        error: format!("worker panicked: {join_err}"),
                    };
                }
            };

        if initial.skip {
            info!(job_id = %job.id, "thumbnails already exist; skipping");
            return JobExecution::Success;
        }

        let bytes = match fetch_version_object(
            &initial.version,
            &storage,
            &initial.version.s3_key,
            state.config.worker_max_document_bytes,
        )
        .await
        {
            Ok(bytes) => bytes,
            Err(err) => return handle_fetch_error(&job, err, "thumbnail fetch failed; will retry"),
        };

        let generation = match generate_preview_and_thumbnail(&initial.document, &bytes) {
            Ok(result) => result,
            Err(err) => {
                return JobExecution::Failed { error: err };
            }
        };

        if let Some(page_count) = generation.page_count {
            let state_clone = state.clone();
            let document_id = initial.document.id;
            let version_id = initial.version.id;
            match task::spawn_blocking(move || {
                persist_document_page_count(state_clone, document_id, version_id, page_count)
            })
            .await
            {
                Ok(Ok(())) => {}
                Ok(Err(err)) => {
                    warn!(
                        job_id = %job.id,
                        document_id = %document_id,
                        version_id = %version_id,
                        error = %err,
                        "failed to update document page count metadata; retrying"
                    );
                    return JobExecution::Retry {
                        delay: Duration::from_secs(30),
                        error: err,
                    };
                }
                Err(join_err) => {
                    error!(
                        job_id = %job.id,
                        document_id = %document_id,
                        version_id = %version_id,
                        error = %join_err,
                        "page count metadata task panicked"
                    );
                    return JobExecution::Retry {
                        delay: Duration::from_secs(60),
                        error: format!("metadata panic: {join_err}"),
                    };
                }
            }
        }

        if let Some(existing_preview) = &initial.existing_preview {
            for object in &initial.existing_preview_objects {
                if let Err(err) = storage.delete_object(&object.s3_key).await {
                    warn!(
                        job_id = %job.id,
                        error = %err,
                        s3_key = %object.s3_key,
                        "failed to delete existing preview object"
                    );
                }
            }

            let tenant_id = initial.document.tenant_id;
            let asset_id = existing_preview.id;
            let state_clone = state.clone();
            match task::spawn_blocking(move || {
                delete_asset(state_clone.as_ref(), tenant_id, asset_id)
            })
            .await
            {
                Ok(Ok(())) => {}
                Ok(Err(err)) => {
                    warn!(
                        job_id = %job.id,
                        error = ?err,
                        asset_id = %asset_id,
                        "failed to remove preview metadata after deletion"
                    );
                }
                Err(join_err) => {
                    warn!(
                        job_id = %job.id,
                        error = %join_err,
                        asset_id = %asset_id,
                        "failed to remove preview metadata: task panicked"
                    );
                }
            }
        }

        if let Some(existing_thumbnail) = &initial.existing_thumbnail {
            for object in &initial.existing_thumbnail_objects {
                if let Err(err) = storage.delete_object(&object.s3_key).await {
                    warn!(
                        job_id = %job.id,
                        error = %err,
                        s3_key = %object.s3_key,
                        "failed to delete existing thumbnail object"
                    );
                }
            }

            let tenant_id = initial.document.tenant_id;
            let asset_id = existing_thumbnail.id;
            let state_clone = state.clone();
            match task::spawn_blocking(move || {
                delete_asset(state_clone.as_ref(), tenant_id, asset_id)
            })
            .await
            {
                Ok(Ok(())) => {}
                Ok(Err(err)) => {
                    warn!(
                        job_id = %job.id,
                        error = ?err,
                        asset_id = %asset_id,
                        "failed to remove thumbnail metadata after deletion"
                    );
                }
                Err(join_err) => {
                    warn!(
                        job_id = %job.id,
                        error = %join_err,
                        asset_id = %asset_id,
                        "failed to remove thumbnail metadata: task panicked"
                    );
                }
            }
        }

        let preview_asset_id = Uuid::new_v4();

        let thumbnail_asset_id = Uuid::new_v4();

        let mut preview_objects: Vec<AssetObjectPersistence> =
            Vec::with_capacity(generation.preview.objects.len());
        for (index, image) in generation.preview.objects.iter().enumerate() {
            if index + 1 > i32::MAX as usize {
                return JobExecution::Failed {
                    error: "too many preview objects".to_string(),
                };
            }
            let ordinal = (index + 1) as i32;
            let s3_key = document_asset_object_key(
                initial.document.id,
                initial.version.version_number,
                PREVIEW_ASSET_TYPE,
                preview_asset_id,
                ordinal,
            );

            if let Err(err) = storage
                .put_object(
                    &s3_key,
                    image.image_bytes.clone(),
                    Some("image/png".into()),
                    None,
                )
                .await
            {
                warn!(job_id = %job.id, error = %err, ordinal, "failed to upload preview; retrying");
                return JobExecution::Retry {
                    delay: Duration::from_secs(30),
                    error: err.to_string(),
                };
            }

            preview_objects.push(AssetObjectPersistence {
                ordinal,
                s3_key,
                width: image.width,
                height: image.height,
            });
        }

        let mut thumbnail_objects: Vec<AssetObjectPersistence> =
            Vec::with_capacity(generation.thumbnail.objects.len());
        for (index, image) in generation.thumbnail.objects.iter().enumerate() {
            if index + 1 > i32::MAX as usize {
                return JobExecution::Failed {
                    error: "too many thumbnail objects".to_string(),
                };
            }
            let ordinal = (index + 1) as i32;
            let s3_key = document_asset_object_key(
                initial.document.id,
                initial.version.version_number,
                THUMBNAIL_ASSET_TYPE,
                thumbnail_asset_id,
                ordinal,
            );

            if let Err(err) = storage
                .put_object(
                    &s3_key,
                    image.image_bytes.clone(),
                    Some("image/png".into()),
                    None,
                )
                .await
            {
                warn!(job_id = %job.id, error = %err, ordinal, "failed to upload thumbnail; retrying");
                return JobExecution::Retry {
                    delay: Duration::from_secs(30),
                    error: err.to_string(),
                };
            }

            thumbnail_objects.push(AssetObjectPersistence {
                ordinal,
                s3_key,
                width: image.width,
                height: image.height,
            });
        }

        let asset_persistences = vec![
            AssetPersistence {
                asset_type: PREVIEW_ASSET_TYPE,
                asset_id: preview_asset_id,
                objects: preview_objects,
            },
            AssetPersistence {
                asset_type: THUMBNAIL_ASSET_TYPE,
                asset_id: thumbnail_asset_id,
                objects: thumbnail_objects,
            },
        ];

        let state_clone = state.clone();
        match task::spawn_blocking(move || {
            persist_assets_metadata(state_clone, &initial, &asset_persistences)
        })
        .await
        {
            Ok(Ok(())) => {}
            Ok(Err(err)) => {
                warn!(job_id = %job.id, error = %err, "failed to persist thumbnail metadata; retrying");
                return JobExecution::Retry {
                    delay: Duration::from_secs(30),
                    error: err,
                };
            }
            Err(join_err) => {
                error!(job_id = %job.id, error = %join_err, "thumbnail metadata update panicked");
                return JobExecution::Retry {
                    delay: Duration::from_secs(30),
                    error: format!("metadata update panic: {join_err}"),
                };
            }
        }

        JobExecution::Success
    }
}

struct ThumbnailContext {
    document: Document,
    version: DocumentVersion,
    existing_thumbnail: Option<DocumentAsset>,
    existing_thumbnail_objects: Vec<DocumentAssetObject>,
    existing_preview: Option<DocumentAsset>,
    existing_preview_objects: Vec<DocumentAssetObject>,
    skip: bool,
}

struct GeneratedImage {
    image_bytes: Vec<u8>,
    width: Option<i32>,
    height: Option<i32>,
}

struct GeneratedAsset {
    objects: Vec<GeneratedImage>,
}

struct GeneratedAssets {
    thumbnail: GeneratedAsset,
    preview: GeneratedAsset,
    page_count: Option<u32>,
}

struct AssetObjectPersistence {
    ordinal: i32,
    s3_key: String,
    width: Option<i32>,
    height: Option<i32>,
}

struct AssetPersistence {
    asset_type: &'static str,
    asset_id: Uuid,
    objects: Vec<AssetObjectPersistence>,
}

fn load_thumbnail_context(
    state: Arc<AppState>,
    payload: &ThumbnailPayload,
) -> Result<ThumbnailContext, String> {
    let mut conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;

    let version: DocumentVersion = document_versions::table
        .find(payload.document_version_id)
        .first(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

    if version.document_id != payload.document_id {
        return Err("document/version mismatch".into());
    }

    let document: Document = documents::table
        .find(payload.document_id)
        .first(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

    let tenant_id = document.tenant_id;

    let existing_assets: Vec<DocumentAsset> = document_assets::table
        .filter(document_assets::document_version_id.eq(payload.document_version_id))
        .filter(document_assets::asset_type.eq_any(vec![
            THUMBNAIL_ASSET_TYPE.to_string(),
            PREVIEW_ASSET_TYPE.to_string(),
        ]))
        .filter(document_assets::tenant_id.eq(tenant_id))
        .load(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

    let mut existing_thumbnail = None;
    let mut existing_thumbnail_objects: Vec<DocumentAssetObject> = Vec::new();
    let mut existing_preview = None;
    let mut existing_preview_objects: Vec<DocumentAssetObject> = Vec::new();
    for asset in existing_assets {
        match asset.asset_type.as_str() {
            THUMBNAIL_ASSET_TYPE => {
                existing_thumbnail_objects = document_asset_objects::table
                    .filter(document_asset_objects::asset_id.eq(asset.id))
                    .filter(document_asset_objects::tenant_id.eq(tenant_id))
                    .order(document_asset_objects::ordinal.asc())
                    .load(&mut conn)
                    .map_err(|err| format!("{err:?}"))?;
                existing_thumbnail = Some(asset);
            }
            PREVIEW_ASSET_TYPE => {
                existing_preview_objects = document_asset_objects::table
                    .filter(document_asset_objects::asset_id.eq(asset.id))
                    .filter(document_asset_objects::tenant_id.eq(tenant_id))
                    .order(document_asset_objects::ordinal.asc())
                    .load(&mut conn)
                    .map_err(|err| format!("{err:?}"))?;
                existing_preview = Some(asset);
            }
            _ => {}
        }
    }

    let (supported, _) = determine_thumbnail_support(&document);
    if !supported {
        return Err("thumbnail generation not supported for this document".into());
    }

    let expected_cardinality = expected_asset_cardinality(&document, &version);
    let preview_cardinality = existing_preview
        .as_ref()
        .and_then(|asset| asset.cardinality)
        .unwrap_or_else(|| existing_preview_objects.len() as i32);
    let thumbnail_cardinality = existing_thumbnail
        .as_ref()
        .and_then(|asset| asset.cardinality)
        .unwrap_or_else(|| existing_thumbnail_objects.len() as i32);

    let needs_regeneration = preview_cardinality < expected_cardinality
        || thumbnail_cardinality < expected_cardinality
        || (existing_preview_objects.len() as i32) < expected_cardinality
        || (existing_thumbnail_objects.len() as i32) < expected_cardinality;

    let skip = existing_thumbnail.is_some()
        && existing_preview.is_some()
        && !payload.force
        && !needs_regeneration;

    Ok(ThumbnailContext {
        document,
        version,
        existing_thumbnail,
        existing_thumbnail_objects,
        existing_preview,
        existing_preview_objects,
        skip,
    })
}

fn generate_preview_and_thumbnail(
    document: &Document,
    bytes: &[u8],
) -> Result<GeneratedAssets, String> {
    let is_pdf = document_is_pdf(document);

    if is_pdf {
        let pdf_assets = generate_pdf_assets(bytes)?;
        Ok(GeneratedAssets {
            preview: pdf_assets.preview,
            thumbnail: pdf_assets.thumbnail,
            page_count: Some(pdf_assets.page_count),
        })
    } else {
        let (preview, thumbnail) = generate_image_assets(bytes)?;
        Ok(GeneratedAssets {
            preview,
            thumbnail,
            page_count: None,
        })
    }
}

fn generate_image_assets(bytes: &[u8]) -> Result<(GeneratedAsset, GeneratedAsset), String> {
    let reader = ImageReader::new(Cursor::new(bytes))
        .with_guessed_format()
        .map_err(|err| err.to_string())?;
    let image = reader.decode().map_err(|err| err.to_string())?;

    let preview_image = if image.width() > PREVIEW_WIDTH || image.height() > PREVIEW_HEIGHT {
        image.thumbnail(PREVIEW_WIDTH, PREVIEW_HEIGHT)
    } else {
        image.clone()
    };

    let thumbnail_image =
        if preview_image.width() > THUMBNAIL_WIDTH || preview_image.height() > THUMBNAIL_HEIGHT {
            preview_image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
        } else {
            preview_image.clone()
        };

    let preview = encode_dynamic_image(preview_image)?;
    let thumbnail = encode_dynamic_image(thumbnail_image)?;

    Ok((
        GeneratedAsset {
            objects: vec![preview],
        },
        GeneratedAsset {
            objects: vec![thumbnail],
        },
    ))
}

struct PdfGeneratedAssets {
    preview: GeneratedAsset,
    thumbnail: GeneratedAsset,
    page_count: u32,
}

fn generate_pdf_assets(bytes: &[u8]) -> Result<PdfGeneratedAssets, String> {
    let pdfium = panic::catch_unwind(|| Pdfium::default())
        .map_err(|_| "failed to initialize PDFium".to_string())?;

    let document = pdfium
        .load_pdf_from_byte_slice(bytes, None)
        .map_err(|err| format!("load pdf: {err}"))?;

    let pages = document.pages();
    let total_pages = pages.len() as usize;

    let render_config = PdfRenderConfig::new()
        .set_target_width(PREVIEW_WIDTH as i32)
        .set_maximum_height(PREVIEW_HEIGHT as i32)
        .render_form_data(true)
        .rotate_if_landscape(PdfPageRenderRotation::None, true);

    let mut preview_objects: Vec<GeneratedImage> = Vec::with_capacity(total_pages);
    let mut thumbnail_objects: Vec<GeneratedImage> = Vec::with_capacity(total_pages);

    for page_index in 0..total_pages {
        let page = pages
            .get(u16::try_from(page_index).map_err(|_| "page index overflow".to_string())?)
            .map_err(|err| format!("load page {page_index}: {err}"))?;

        let bitmap = page
            .render_with_config(&render_config)
            .map_err(|err| format!("render pdf page {page_index}: {err}"))?;

        let preview_buffer = bitmap.as_image().to_rgb8();
        let preview_image = image::DynamicImage::ImageRgb8(preview_buffer);

        let thumbnail_image = if preview_image.width() > THUMBNAIL_WIDTH
            || preview_image.height() > THUMBNAIL_HEIGHT
        {
            preview_image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
        } else {
            preview_image.clone()
        };

        preview_objects.push(encode_dynamic_image(preview_image)?);
        thumbnail_objects.push(encode_dynamic_image(thumbnail_image)?);
    }

    let page_count: u32 = total_pages
        .try_into()
        .map_err(|_| "page count exceeds supported range".to_string())?;

    Ok(PdfGeneratedAssets {
        preview: GeneratedAsset {
            objects: preview_objects,
        },
        thumbnail: GeneratedAsset {
            objects: thumbnail_objects,
        },
        page_count,
    })
}

fn encode_dynamic_image(image: image::DynamicImage) -> Result<GeneratedImage, String> {
    let (width, height) = image.dimensions();
    let mut cursor = Cursor::new(Vec::new());
    image
        .write_to(&mut cursor, ImageFormat::Png)
        .map_err(|err| err.to_string())?;
    Ok(GeneratedImage {
        image_bytes: cursor.into_inner(),
        width: Some(width as i32),
        height: Some(height as i32),
    })
}

fn persist_assets_metadata(
    state: Arc<AppState>,
    context: &ThumbnailContext,
    assets: &[AssetPersistence],
) -> Result<(), String> {
    let tenant_id = context.document.tenant_id;
    let mut conn = state
        .db_for_tenant(tenant_id)
        .map_err(|err| format!("{err:?}"))?;

    if let Some(existing_preview) = &context.existing_preview {
        diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_preview.id)))
            .execute(&mut conn)
            .map_err(|err| format!("{err:?}"))?;
    }

    if let Some(existing_thumbnail) = &context.existing_thumbnail {
        diesel::delete(
            document_assets::table.filter(document_assets::id.eq(existing_thumbnail.id)),
        )
        .execute(&mut conn)
        .map_err(|err| format!("{err:?}"))?;
    }

    for asset in assets {
        if asset.objects.is_empty() {
            return Err(format!(
                "asset {} has no generated objects",
                asset.asset_type
            ));
        }

        let object_count: i32 = asset
            .objects
            .len()
            .try_into()
            .map_err(|_| "asset contains too many objects".to_string())?;

        let new_asset = NewDocumentAsset {
            id: asset.asset_id,
            document_version_id: context.version.id,
            asset_type: asset.asset_type.to_string(),
            mime_type: "image/png".to_string(),
            metadata: json!({
                "generated_at": Utc::now().to_rfc3339(),
            }),
            cardinality: Some(object_count),
            tenant_id,
        };

        diesel::insert_into(document_assets::table)
            .values(&new_asset)
            .on_conflict((
                document_assets::document_version_id,
                document_assets::asset_type,
            ))
            .do_update()
            .set((
                document_assets::mime_type.eq(excluded(document_assets::mime_type)),
                document_assets::metadata.eq(excluded(document_assets::metadata)),
                document_assets::cardinality.eq(excluded(document_assets::cardinality)),
            ))
            .execute(&mut conn)
            .map_err(|err| format!("{err:?}"))?;

        diesel::delete(
            document_asset_objects::table
                .filter(document_asset_objects::asset_id.eq(asset.asset_id))
                .filter(document_asset_objects::tenant_id.eq(tenant_id)),
        )
        .execute(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

        for object in &asset.objects {
            let mut metadata_map = Map::new();
            if let Some(width) = object.width {
                metadata_map.insert("width".to_string(), Value::from(width));
            }
            if let Some(height) = object.height {
                metadata_map.insert("height".to_string(), Value::from(height));
            }

            let object_metadata = Value::Object(metadata_map);

            let new_object = NewDocumentAssetObject {
                id: Uuid::new_v4(),
                asset_id: asset.asset_id,
                ordinal: object.ordinal,
                s3_key: object.s3_key.clone(),
                metadata: object_metadata,
                tenant_id,
            };

            diesel::insert_into(document_asset_objects::table)
                .values(&new_object)
                .execute(&mut conn)
                .map_err(|err| format!("{err:?}"))?;
        }
    }

    Ok(())
}

fn persist_document_page_count(
    state: Arc<AppState>,
    document_id: Uuid,
    document_version_id: Uuid,
    page_count: u32,
) -> Result<(), String> {
    let mut conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
    let tenant_id: Uuid = documents::table
        .find(document_id)
        .select(documents::tenant_id)
        .first(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

    let existing_metadata: Value = document_versions::table
        .filter(document_versions::id.eq(document_version_id))
        .filter(document_versions::document_id.eq(document_id))
        .filter(document_versions::tenant_id.eq(tenant_id))
        .select(document_versions::metadata)
        .first(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

    let updated = match existing_metadata {
        Value::Object(mut map) => {
            map.insert("page_count".to_string(), Value::from(page_count));
            Value::Object(map)
        }
        _ => {
            let mut map = Map::new();
            map.insert("page_count".to_string(), Value::from(page_count));
            Value::Object(map)
        }
    };

    diesel::update(
        document_versions::table
            .filter(document_versions::id.eq(document_version_id))
            .filter(document_versions::document_id.eq(document_id))
            .filter(document_versions::tenant_id.eq(tenant_id)),
    )
    .set(document_versions::metadata.eq(updated))
    .execute(&mut conn)
    .map_err(|err| format!("{err:?}"))?;

    Ok(())
}

fn document_is_pdf(document: &Document) -> bool {
    document
        .content_type
        .as_deref()
        .map(|mime| mime.eq_ignore_ascii_case("application/pdf"))
        .unwrap_or_else(|| {
            document
                .original_name
                .rsplit('.')
                .next()
                .map(|ext| ext.eq_ignore_ascii_case("pdf"))
                .unwrap_or(false)
        })
}

fn expected_asset_cardinality(document: &Document, version: &DocumentVersion) -> i32 {
    if let Value::Object(map) = &version.metadata {
        if let Some(count) = map.get("page_count").and_then(|v| v.as_i64()) {
            if count > 0 {
                return count
                    .min(i64::from(i32::MAX))
                    .try_into()
                    .unwrap_or(i32::MAX);
            }
        }
    }

    if document_is_pdf(document) {
        1
    } else {
        1
    }
}