pub mod error;
pub use error::OpenAIExtractorError;

use crate::Info;
use async_openai::{
    Client,
    types::responses::{
        CreateResponseArgs, Response, ResponseFormatJsonSchema, ResponseTextParam, Status,
        TextResponseFormatConfiguration,
    },
};
use serde::de::Error as _;
use std::{marker::PhantomData, time::Duration};

use super::{Extract, ExtractorError};

const MAX_RETRIES: u32 = 5;
const MAX_BATCH_INPUT_TOKENS: usize = 100_000;
const SYSTEM_PROMPT_OVERHEAD: usize = 500;
const OUTPUT_TOKENS_RESERVED: usize = 20_000;
const MAX_TOKENS_PER_ITEM: usize = 10_000;
const MAX_ITEMS_PER_BATCH: usize = 12;
const MIN_OUTPUT_TOKENS_PER_BATCH: u32 = 4_000;
const SINGLE_ITEM_MAX_OUTPUT_TOKENS: u32 = 8_000;

fn strict_json_schema(mut schema: serde_json::Value) -> serde_json::Value {
    apply_strict_object_schema(&mut schema);
    schema
}

fn apply_strict_object_schema(schema: &mut serde_json::Value) {
    let Some(object) = schema.as_object_mut() else {
        return;
    };

    if matches!(object.get("type"), Some(serde_json::Value::String(kind)) if kind == "object") {
        object
            .entry("additionalProperties".to_string())
            .or_insert(serde_json::Value::Bool(false));
        if let Some(properties) = object
            .get("properties")
            .and_then(serde_json::Value::as_object)
        {
            let required = properties
                .keys()
                .cloned()
                .map(serde_json::Value::String)
                .collect::<Vec<_>>();
            object.insert("required".to_string(), serde_json::Value::Array(required));
        }
    }

    if let Some(properties) = object
        .get_mut("properties")
        .and_then(serde_json::Value::as_object_mut)
    {
        for value in properties.values_mut() {
            apply_strict_object_schema(value);
        }
    }

    if let Some(items) = object.get_mut("items") {
        apply_strict_object_schema(items);
    }

    if let Some(any_of) = object
        .get_mut("anyOf")
        .and_then(serde_json::Value::as_array_mut)
    {
        for value in any_of {
            apply_strict_object_schema(value);
        }
    }

    if let Some(one_of) = object
        .get_mut("oneOf")
        .and_then(serde_json::Value::as_array_mut)
    {
        for value in one_of {
            apply_strict_object_schema(value);
        }
    }

    if let Some(all_of) = object
        .get_mut("allOf")
        .and_then(serde_json::Value::as_array_mut)
    {
        for value in all_of {
            apply_strict_object_schema(value);
        }
    }
}

fn response_format_name(schema: &serde_json::Value) -> String {
    schema
        .get("title")
        .and_then(serde_json::Value::as_str)
        .map(sanitize_response_format_name)
        .filter(|name| !name.is_empty())
        .unwrap_or_else(|| "info_array".to_string())
}

fn sanitize_response_format_name(name: &str) -> String {
    let sanitized = name
        .chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '_' || *c == '-')
        .take(64)
        .collect::<String>();
    if sanitized.is_empty() {
        "info_array".to_string()
    } else {
        sanitized
    }
}

fn response_format_description(schema: &serde_json::Value) -> String {
    schema
        .get("description")
        .and_then(serde_json::Value::as_str)
        .map(str::trim)
        .filter(|description| !description.is_empty())
        .unwrap_or("Array of extracted information")
        .to_string()
}

fn append_prompt_context<T: Info>(prompt: &mut String, prior_items: &[T]) {
    if let Some(extra) = T::extraction_instructions() {
        prompt.push_str("\n\nAdditional extraction rules:\n");
        prompt.push_str(extra);
    }

    if let Some(memory) = T::continuity_instructions(prior_items) {
        let memory = memory.trim();
        if !memory.is_empty() {
            prompt.push_str("\n\nCanonical outputs already accepted in this run:\n");
            prompt.push_str(memory);
        }
    }
}

fn build_response_format<T: Info>() -> ResponseFormatJsonSchema {
    let item_schema = strict_json_schema(<T as crate::Info>::json_schema());

    ResponseFormatJsonSchema {
        name: response_format_name(&item_schema),
        description: Some(response_format_description(&item_schema)),
        schema: Some(serde_json::json!({
            "type": "object",
            "properties": {
                "items": {
                    "type": "array",
                    "items": item_schema
                }
            },
            "required": ["items"],
            "additionalProperties": false
        })),
        strict: Some(true),
    }
}

fn response_was_truncated(response: &Response) -> bool {
    response.status == Status::Incomplete
        && response
            .incomplete_details
            .as_ref()
            .map(|details| details.reason == "max_output_tokens")
            .unwrap_or(false)
}

fn available_batch_input_tokens() -> usize {
    MAX_BATCH_INPUT_TOKENS
        .saturating_sub(SYSTEM_PROMPT_OVERHEAD)
        .saturating_sub(OUTPUT_TOKENS_RESERVED)
}

fn max_output_tokens_for_batch(items_len: usize) -> u32 {
    if items_len == 1 {
        return SINGLE_ITEM_MAX_OUTPUT_TOKENS;
    }

    ((items_len as u32) * 1_500)
        .max(MIN_OUTPUT_TOKENS_PER_BATCH)
        .min(32_000)
}

fn build_batches(item_tokens: &[usize]) -> Vec<Vec<usize>> {
    let available_tokens = available_batch_input_tokens();
    let mut batches = Vec::new();
    let mut current_batch = Vec::new();
    let mut current_tokens = 0;

    for (idx, &tokens) in item_tokens.iter().enumerate() {
        let capped_tokens = tokens.min(MAX_TOKENS_PER_ITEM);
        if capped_tokens > available_tokens {
            flush_batch(&mut batches, &mut current_batch, &mut current_tokens);
            batches.push(vec![idx]);
            continue;
        }

        let would_exceed_tokens = current_tokens + capped_tokens > available_tokens;
        let would_exceed_items = current_batch.len() >= MAX_ITEMS_PER_BATCH;
        if would_exceed_tokens || would_exceed_items {
            flush_batch(&mut batches, &mut current_batch, &mut current_tokens);
        }

        current_batch.push(idx);
        current_tokens += capped_tokens;
    }

    flush_batch(&mut batches, &mut current_batch, &mut current_tokens);
    batches
}

fn estimate_tokens(text: &str) -> usize {
    text.len() / 4
}

fn truncate_to_tokens(text: &str, max_tokens: usize) -> String {
    let max_chars = max_tokens * 4;
    if text.len() <= max_chars {
        text.to_string()
    } else {
        let safe_boundary = text.floor_char_boundary(max_chars);
        format!("{}... [truncated]", &text[..safe_boundary])
    }
}

fn is_rate_limit_error(e: &OpenAIExtractorError) -> bool {
    let msg = e.to_string().to_lowercase();
    msg.contains("rate limit") || msg.contains("rate_limit") || msg.contains("429")
}

fn is_oversized_rate_limit_error(e: &OpenAIExtractorError) -> bool {
    let msg = e.to_string().to_lowercase();
    msg.contains("request too large")
        || msg.contains("tokens per min")
        || (is_rate_limit_error(e) && msg.contains("tpm"))
}

/// Parse the suggested retry delay from an OpenAI rate-limit message.
/// Handles "Xm Y.Zs", "XmY.Zs", "X.Zs", "Xs". Returns seconds.
fn parse_retry_after(e: &OpenAIExtractorError) -> f64 {
    let msg = e.to_string();
    if let Some(caps) = regex::Regex::new(r"(\d+)m(\d+(?:\.\d+)?)s")
        .ok()
        .and_then(|re| re.captures(&msg))
    {
        let mins: f64 = caps[1].parse().unwrap_or(0.0);
        let secs: f64 = caps[2].parse().unwrap_or(0.0);
        return mins * 60.0 + secs;
    }
    if let Some(caps) = regex::Regex::new(r"(\d+(?:\.\d+)?)s")
        .ok()
        .and_then(|re| re.captures(&msg))
        && let Ok(secs) = caps[1].parse::<f64>()
    {
        return secs;
    }
    60.0
}

/// Default OpenAI model used by [`OpenAIExtractor::new`].
pub const DEFAULT_MODEL: &str = "gpt-5-nano";

/// OpenAI-backed extractor for `Info` records.
#[derive(Clone)]
pub struct OpenAIExtractor<T: Info> {
    client: Client<async_openai::config::OpenAIConfig>,
    model: String,
    _phantom: PhantomData<T>,
}

impl<T: Info + Clone + Send + Sync> OpenAIExtractor<T> {
    /// Build an extractor with [`DEFAULT_MODEL`].
    pub fn new() -> Result<Self, ExtractorError> {
        Self::with_model(DEFAULT_MODEL)
    }

    /// Build an extractor with a specific model name.
    pub fn with_model(model: impl Into<String>) -> Result<Self, ExtractorError> {
        let api_key = option_env!("OPENAI_API_KEY")
            .map(|s| s.to_string())
            .or_else(|| std::env::var("OPENAI_API_KEY").ok())
            .ok_or(OpenAIExtractorError::MissingApiKey)?;
        let config = async_openai::config::OpenAIConfig::new().with_api_key(&api_key);
        Ok(Self {
            client: Client::with_config(config),
            model: model.into(),
            _phantom: PhantomData,
        })
    }

    #[async_recursion::async_recursion]
    async fn extract_with_split_and_retry(
        &self,
        items: &[(String, Vec<String>)],
        prior_items: &[T],
    ) -> Result<Vec<(String, T)>, ExtractorError> {
        let mut attempt = 0u32;
        loop {
            match self.extract_single_batch(items, prior_items).await {
                Ok(r) => return Ok(r),

                Err(OpenAIExtractorError::OutputTruncated) if items.len() > 1 => {
                    return self.split_batch_and_retry(items, prior_items, None).await;
                }

                Err(e) if is_oversized_rate_limit_error(&e) && items.len() > 1 => {
                    return self
                        .split_batch_and_retry(items, prior_items, Some(&e))
                        .await;
                }

                Err(e) if is_rate_limit_error(&e) && attempt < MAX_RETRIES => {
                    let wait = parse_retry_after(&e) + 2.0;
                    log::warn!(
                        "⏳ Rate limited — waiting {:.0}s then retrying (attempt {}/{})",
                        wait,
                        attempt + 1,
                        MAX_RETRIES
                    );
                    tokio::time::sleep(Duration::from_secs_f64(wait)).await;
                    attempt += 1;
                }

                Err(e) if items.len() > 1 => {
                    return self
                        .split_batch_and_retry(items, prior_items, Some(&e))
                        .await;
                }

                Err(e) => return Err(e.into()),
            }
        }
    }

    async fn split_batch_and_retry(
        &self,
        items: &[(String, Vec<String>)],
        prior_items: &[T],
        error: Option<&OpenAIExtractorError>,
    ) -> Result<Vec<(String, T)>, ExtractorError> {
        let mid = items.len() / 2;
        if let Some(error) = error {
            log::warn!(
                "  Batch of {} failed ({}); splitting into halves of {} and {}",
                items.len(),
                error,
                mid,
                items.len() - mid
            );
        } else {
            log::info!(
                "  Splitting batch of {} into halves of {} and {}",
                items.len(),
                mid,
                items.len() - mid
            );
        }

        let mut left = self
            .extract_with_split_and_retry(&items[..mid], prior_items)
            .await?;
        let mut right_prior_items = prior_items.to_vec();
        right_prior_items.extend(left.iter().map(|(_, info)| info.clone()));
        let right = self
            .extract_with_split_and_retry(&items[mid..], &right_prior_items)
            .await?;
        left.extend(right);
        Ok(left)
    }

    async fn extract_single_batch(
        &self,
        items: &[(String, Vec<String>)],
        prior_items: &[T],
    ) -> Result<Vec<(String, T)>, OpenAIExtractorError> {
        let combined_text = items
            .iter()
            .enumerate()
            .map(|(idx, (key, texts))| {
                let pages = texts
                    .iter()
                    .enumerate()
                    .map(|(page_idx, text)| format!("  Page {}: {}\n", page_idx + 1, text))
                    .collect::<Vec<_>>()
                    .join("\n");

                let item_text = format!("=== Item {} - {} ===\n{}\n", idx + 1, key, pages);
                let original_tokens = estimate_tokens(&item_text);
                let truncated = truncate_to_tokens(&item_text, MAX_TOKENS_PER_ITEM);
                if original_tokens > MAX_TOKENS_PER_ITEM {
                    log::debug!(
                        "  Item {} ({}) truncated from ~{} to ~{} tokens",
                        idx + 1,
                        key,
                        original_tokens,
                        MAX_TOKENS_PER_ITEM
                    );
                }
                truncated
            })
            .collect::<Vec<_>>()
            .join("\n");

        log::debug!(
            "Batch text length: {} chars for {} items",
            combined_text.len(),
            items.len()
        );
        log::debug!(
            "Estimated input tokens: ~{}",
            estimate_tokens(&combined_text)
        );
        log::debug!("Building OpenAI request for batch of {} items", items.len());

        let mut system_prompt = format!(
            "You are a data extraction assistant. Information from {} different sources is provided. Extract information for EACH source and return it in the specified JSON array format. Process ALL {} sources in a SINGLE response, one object per source in the exact order provided. Combine information from all pages for each source.",
            items.len(),
            items.len()
        );
        append_prompt_context::<T>(&mut system_prompt, prior_items);

        let user_prompt = format!("Extract information for each item:\n\n{}", combined_text);

        let max_output_tokens = max_output_tokens_for_batch(items.len());
        let mut request = CreateResponseArgs::default();
        request.model(&self.model);
        request.input(user_prompt);
        request.instructions(system_prompt);
        request.max_output_tokens(max_output_tokens);
        request.text(ResponseTextParam {
            format: TextResponseFormatConfiguration::JsonSchema(build_response_format::<T>()),
            verbosity: None,
        });
        let request = request.build()?;

        log::debug!(
            "Sending OpenAI request for batch of {} items (~{} input tokens, max {} output tokens)",
            items.len(),
            estimate_tokens(&combined_text),
            max_output_tokens
        );
        let response = self.client.responses().create(request).await?;
        log::debug!(
            "Received OpenAI response for batch of {} items",
            items.len()
        );

        if response_was_truncated(&response) {
            log::warn!(
                "⚠️  Output truncated (status: incomplete/max_output_tokens) for batch of {} items — will split",
                items.len()
            );
            return Err(OpenAIExtractorError::OutputTruncated);
        }

        let content = response
            .output_text()
            .map(|text| text.trim().to_string())
            .filter(|text| !text.is_empty())
            .ok_or_else(|| {
                log::error!(
                    "❌ Empty response from OpenAI — status: {:?}",
                    response.status
                );
                OpenAIExtractorError::SerdeJsonError(serde_json::Error::custom(format!(
                    "Empty response content (status: {:?})",
                    response.status
                )))
            })?;

        let wrapper: serde_json::Value = serde_json::from_str(&content)?;
        let items_array = wrapper.get("items").ok_or_else(|| {
            OpenAIExtractorError::SerdeJsonError(serde_json::Error::custom(
                "Missing 'items' field in response",
            ))
        })?;
        let mut infos: Vec<T> = serde_json::from_value(items_array.clone())?;

        let mut results = Vec::new();
        for (info, (key, texts)) in infos.iter_mut().zip(items.iter()) {
            let combined_text = texts.join("\n");
            info.post_process(&combined_text).map_err(|e| {
                OpenAIExtractorError::SerdeJsonError(serde_json::Error::custom(e.to_string()))
            })?;
            results.push((key.clone(), info.clone()));
        }

        Ok(results)
    }
}

impl<T: Info + Clone + Send + Sync + 'static> Extract<T> for OpenAIExtractor<T> {
    async fn extract(
        &self,
        items: &[(String, Vec<String>)],
        on_batch: &(dyn Fn(Vec<(String, T)>) + Send + Sync),
    ) -> Result<Vec<(String, T)>, ExtractorError> {
        log::info!(
            "Batch extracting from {} items with automatic batching",
            items.len()
        );

        let item_tokens: Vec<usize> = items
            .iter()
            .map(|(key, texts)| {
                let combined = texts.join("\n");
                let tokens = estimate_tokens(&format!("=== Item - {} ===\n{}\n", key, combined));
                tokens.min(MAX_TOKENS_PER_ITEM)
            })
            .collect();

        let available_tokens = available_batch_input_tokens();
        for (idx, &tokens) in item_tokens.iter().enumerate() {
            if tokens > available_tokens {
                log::warn!(
                    "Item {} ({}) has {} estimated tokens, exceeding limit of {}. Processing alone with truncation.",
                    idx + 1,
                    items[idx].0,
                    tokens,
                    available_tokens
                );
            }
        }

        let batches = build_batches(&item_tokens);

        log::info!("Split {} items into {} batches", items.len(), batches.len());
        for (batch_idx, batch) in batches.iter().enumerate() {
            let total_tokens: usize = batch.iter().map(|&i| item_tokens[i]).sum();
            log::info!(
                "  Batch {}/{}: {} items, ~{} tokens",
                batch_idx + 1,
                batches.len(),
                batch.len(),
                total_tokens
            );
        }

        let total_batches = batches.len();
        let mut all_results = Vec::new();
        let mut prior_items = Vec::new();

        for (batch_idx, batch_indices) in batches.iter().enumerate() {
            let batch_items: Vec<_> = batch_indices
                .iter()
                .map(|&idx| items[idx].clone())
                .collect();

            log::info!(
                "Processing batch {}/{} ({} items)...",
                batch_idx + 1,
                total_batches,
                batch_items.len()
            );

            let results = self
                .extract_with_split_and_retry(&batch_items, &prior_items)
                .await?;
            log::info!("  ✅ Batch {}/{} completed", batch_idx + 1, total_batches);
            prior_items.extend(results.iter().map(|(_, info)| info.clone()));
            on_batch(results.clone());
            all_results.extend(results);
        }

        Ok(all_results)
    }
}

fn flush_batch(
    batches: &mut Vec<Vec<usize>>,
    current_batch: &mut Vec<usize>,
    current_tokens: &mut usize,
) {
    if !current_batch.is_empty() {
        batches.push(std::mem::take(current_batch));
        *current_tokens = 0;
    }
}

#[cfg(test)]
mod tests;
