use std::{
    fmt, fs,
    io::{ErrorKind, Write},
    process::Command,
    sync::Arc,
    time::Duration,
};

use async_trait::async_trait;
use chrono::Utc;
use diesel::{pg::upsert::excluded, prelude::*};
use pdfium_render::prelude::*;
use serde::Deserialize;
use serde_json::json;
use tempfile::NamedTempFile;
use tokio::task;
use tracing::{error, info, warn};
use uuid::Uuid;

use crate::{
    documents::asset::delete_asset,
    jobs::{enqueue_job, JOB_GENERATE_OCR_TEXT, JOB_INDEX_DOCUMENT_TEXT},
    models::{
        Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
        NewDocumentAssetObject,
    },
    schema::{document_asset_objects, document_assets},
    state::AppState,
    storage::TenantStorage,
    utils::storage_paths::document_asset_object_prefix,
};

use super::{
    common::{load_document_version, load_version_assets},
    fetch_version_object, handle_fetch_error, JobExecution, JobHandler,
};

pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text";
const MIN_TEXT_LENGTH: usize = 50;

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

pub struct GenerateOcrTextJob;

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

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

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

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

        if context.skip {
            info!(job_id = %job.id, "ocr already present; skipping");
            return JobExecution::Success;
        }

        let bytes = match fetch_version_object(
            &context.version,
            &storage,
            &context.version.s3_key,
            state.config.worker_max_document_bytes,
        )
        .await
        {
            Ok(bytes) => bytes,
            Err(err) => return handle_fetch_error(&job, err, "failed to fetch document for ocr"),
        };

        let doc_meta = PdfDocumentMeta {
            content_type: context.document.content_type.clone(),
            original_name: context.document.original_name.clone(),
        };

        let generation =
            match task::spawn_blocking(move || generate_ocr_text(&doc_meta, &bytes)).await {
                Ok(result) => result,
                Err(join_err) => {
                    error!(job_id = %job.id, error = %join_err, "ocr text task panicked");
                    return JobExecution::Retry {
                        delay: Duration::from_secs(60),
                        error: format!("worker panicked: {join_err}"),
                    };
                }
            };

        let Some(generation) = generation else {
            warn!(job_id = %job.id, "no text extracted from document; failing job");
            return JobExecution::Failed {
                error: "no text extracted and OCR unavailable".into(),
            };
        };

        if let Some(existing_asset) = &context.existing_asset {
            for object in &context.existing_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 ocr asset object");
                }
            }

            let tenant_id = context.document.tenant_id;
            let asset_id = existing_asset.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 ocr asset metadata after deletion");
                }
                Err(join_err) => {
                    warn!(job_id = %job.id, error = %join_err, asset_id = %asset_id, "failed to remove ocr asset metadata: task panicked");
                }
            }
        }

        let asset_id = Uuid::new_v4();

        let s3_key = document_asset_object_prefix(
            context.document.id,
            context.version.version_number,
            OCR_TEXT_ASSET_TYPE,
            asset_id,
        );

        if let Err(err) = storage
            .put_object(
                &s3_key,
                generation.text.into_bytes(),
                Some("text/plain; charset=utf-8".into()),
                None,
            )
            .await
        {
            warn!(job_id = %job.id, error = %err, "failed to upload ocr text");
            return JobExecution::Retry {
                delay: Duration::from_secs(30),
                error: err.to_string(),
            };
        }

        let state_clone = state.clone();
        match task::spawn_blocking(move || {
            persist_ocr_metadata(state_clone, &context, asset_id, &s3_key, generation.source)
        })
        .await
        {
            Ok(Ok(())) => {
                if let Err(err) = enqueue_index_job(&state, job.tenant_id, &payload) {
                    warn!(job_id = %job.id, error = %err, "failed to enqueue index job");
                }
                JobExecution::Success
            }
            Ok(Err(err)) => {
                warn!(job_id = %job.id, error = %err, "failed to persist ocr metadata");
                JobExecution::Retry {
                    delay: Duration::from_secs(30),
                    error: err,
                }
            }
            Err(join_err) => {
                error!(job_id = %job.id, error = %join_err, "ocr metadata task panicked");
                JobExecution::Retry {
                    delay: Duration::from_secs(30),
                    error: format!("metadata update panic: {join_err}"),
                }
            }
        }
    }
}

struct PdfDocumentMeta {
    content_type: Option<String>,
    original_name: String,
}

struct OcrContext {
    document: Document,
    version: DocumentVersion,
    existing_asset: Option<DocumentAsset>,
    existing_objects: Vec<DocumentAssetObject>,
    skip: bool,
}

struct OcrGeneration {
    text: String,
    source: &'static str,
}

fn load_ocr_context(
    state: Arc<AppState>,
    tenant_id: Uuid,
    payload: OcrPayload,
) -> Result<OcrContext, String> {
    let base = load_document_version(
        state.as_ref(),
        tenant_id,
        payload.document_id,
        payload.document_version_id,
    )?;

    let mut conn = state
        .db_for_tenant(base.tenant_id)
        .map_err(|err| format!("{err:?}"))?;

    let mut assets = load_version_assets(
        &mut conn,
        base.tenant_id,
        base.version.id,
        &[OCR_TEXT_ASSET_TYPE],
    )?;

    let (existing_asset, existing_objects) = assets
        .remove(OCR_TEXT_ASSET_TYPE)
        .map(|entry| (Some(entry.asset), entry.objects))
        .unwrap_or((None, Vec::new()));

    let is_pdf = document_is_pdf(&base.document);
    if !is_pdf {
        return Ok(OcrContext {
            document: base.document,
            version: base.version,
            existing_asset,
            existing_objects,
            skip: true,
        });
    }

    let skip = existing_asset.is_some() && !payload.force;

    Ok(OcrContext {
        document: base.document,
        version: base.version,
        existing_asset,
        existing_objects,
        skip,
    })
}

fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option<OcrGeneration> {
    if !document_meta_is_pdf(meta) {
        return None;
    }

    if let Ok(text) = extract_pdf_text(bytes) {
        if text.trim().chars().count() >= MIN_TEXT_LENGTH {
            return Some(OcrGeneration {
                text,
                source: "pdf-text",
            });
        }
    }

    match run_ocr(bytes) {
        Ok(Some(text)) => Some(OcrGeneration {
            text,
            source: "ocr",
        }),
        Ok(None) => None,
        Err(OcrError::BinaryMissing) => {
            warn!("ocrmypdf not installed; cannot perform OCR");
            None
        }
        Err(err) => {
            warn!(error = ?err, "ocr command failed");
            None
        }
    }
}

fn extract_pdf_text(bytes: &[u8]) -> Result<String, String> {
    let pdfium = Pdfium::default();
    let document = pdfium
        .load_pdf_from_byte_slice(bytes, None)
        .map_err(|err| format!("load pdf: {err}"))?;

    let mut combined = String::new();
    let pages = document.pages();
    for page_index in 0..pages.len() {
        let page = pages
            .get(page_index)
            .map_err(|err| format!("load page {page_index}: {err}"))?;
        if let Ok(page_text) = page.text() {
            for segment in page_text.segments().iter() {
                combined.push_str(&segment.text());
                combined.push('\n');
            }
        };
    }

    Ok(combined)
}

#[derive(Debug)]
enum OcrError {
    BinaryMissing,
    Failed(String),
}

impl fmt::Display for OcrError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            OcrError::BinaryMissing => write!(f, "ocrmypdf binary not found"),
            OcrError::Failed(msg) => write!(f, "ocr failed: {msg}"),
        }
    }
}

fn run_ocr(bytes: &[u8]) -> Result<Option<String>, OcrError> {
    let mut input = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
    input
        .write_all(bytes)
        .map_err(|err| OcrError::Failed(err.to_string()))?;
    input
        .flush()
        .map_err(|err| OcrError::Failed(err.to_string()))?;

    let output_pdf = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
    let sidecar = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;

    let status = Command::new("ocrmypdf")
        .arg("--sidecar")
        .arg(sidecar.path())
        .arg("--skip-text")
        .arg(input.path())
        .arg(output_pdf.path())
        .output();

    match status {
        Ok(output) => {
            if !output.status.success() {
                return Err(OcrError::Failed(format!(
                    "ocrmypdf failed: exit={} stderr={}",
                    output.status,
                    String::from_utf8_lossy(&output.stderr)
                )));
            }

            let text = fs::read_to_string(sidecar.path())
                .map_err(|err| OcrError::Failed(err.to_string()))?;
            if text.trim().chars().count() >= MIN_TEXT_LENGTH {
                Ok(Some(text))
            } else {
                Ok(None)
            }
        }
        Err(err) => {
            if err.kind() == ErrorKind::NotFound {
                Err(OcrError::BinaryMissing)
            } else {
                Err(OcrError::Failed(err.to_string()))
            }
        }
    }
}

fn persist_ocr_metadata(
    state: Arc<AppState>,
    context: &OcrContext,
    asset_id: Uuid,
    s3_key: &str,
    source: &'static str,
) -> 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_asset) = &context.existing_asset {
        diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_asset.id)))
            .execute(&mut conn)
            .map_err(|err| format!("{err:?}"))?;
    }

    let new_asset = NewDocumentAsset {
        id: asset_id,
        document_version_id: context.version.id,
        asset_type: OCR_TEXT_ASSET_TYPE.to_string(),
        mime_type: "text/plain".to_string(),
        metadata: json!({
            "generated_at": Utc::now().to_rfc3339(),
            "source": source,
        }),
        cardinality: Some(1),
        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:?}"))?;

    let existing_object_id: Option<Uuid> = document_asset_objects::table
        .filter(document_asset_objects::asset_id.eq(asset_id))
        .filter(document_asset_objects::ordinal.eq(1))
        .filter(document_asset_objects::tenant_id.eq(tenant_id))
        .select(document_asset_objects::id)
        .first(&mut conn)
        .optional()
        .map_err(|err| format!("{err:?}"))?;

    let object_id = existing_object_id.unwrap_or_else(Uuid::new_v4);

    let new_object = NewDocumentAssetObject {
        id: object_id,
        asset_id,
        ordinal: 1,
        s3_key: s3_key.to_string(),
        metadata: json!({}),
        tenant_id,
    };

    diesel::insert_into(document_asset_objects::table)
        .values(&new_object)
        .on_conflict((
            document_asset_objects::asset_id,
            document_asset_objects::ordinal,
        ))
        .do_update()
        .set((
            document_asset_objects::s3_key.eq(excluded(document_asset_objects::s3_key)),
            document_asset_objects::metadata.eq(excluded(document_asset_objects::metadata)),
        ))
        .execute(&mut conn)
        .map_err(|err| format!("{err:?}"))?;

    Ok(())
}

fn enqueue_index_job(
    state: &AppState,
    tenant_id: Uuid,
    payload: &OcrPayload,
) -> Result<(), String> {
    let mut conn = state
        .db_for_tenant(tenant_id)
        .map_err(|err| format!("{err:?}"))?;
    enqueue_job(
        &mut conn,
        tenant_id,
        JOB_INDEX_DOCUMENT_TEXT,
        json!({
            "document_id": payload.document_id,
            "document_version_id": payload.document_version_id,
        }),
        None,
    )
    .map(|_| ())
    .map_err(|err| err.to_string())
}

pub fn document_is_pdf(document: &Document) -> bool {
    document_meta_is_pdf(&PdfDocumentMeta {
        content_type: document.content_type.clone(),
        original_name: document.original_name.clone(),
    })
}

fn document_meta_is_pdf(meta: &PdfDocumentMeta) -> bool {
    if let Some(content_type) = &meta.content_type {
        if content_type.eq_ignore_ascii_case("application/pdf") {
            return true;
        }
    }

    meta.original_name
        .rsplit('.')
        .next()
        .map(|ext| ext.eq_ignore_ascii_case("pdf"))
        .unwrap_or(false)
}
