sds-converter-core 0.1.2

use std::time::Duration;

use reqwest::Client;
use serde_json::Value;

use crate::error::SdsError;
use crate::language::Language;
use crate::schema::SdsRoot;

// ---------------------------------------------------------------------------
// LlmBackend trait (stable async fn in trait since Rust 1.75)
// ---------------------------------------------------------------------------

/// Abstraction over LLM completion providers.
///
/// Implement this trait to use any LLM backend with `sds-converter-core`.
/// The library ships with [`AnthropicBackend`] and [`OpenAiCompatBackend`].
pub trait LlmBackend {
    /// Send a system + user message pair and return the raw text response.
    fn complete(
        &self,
        system: &str,
        user: &str,
    ) -> impl std::future::Future<Output = Result<String, SdsError>> + Send;
}

// ---------------------------------------------------------------------------
// Provider URL table
// ---------------------------------------------------------------------------

/// Returns the default base URL for well-known OpenAI-compatible providers.
///
/// Recognised names: `"openai"`, `"gemini"`, `"mistral"`, `"groq"`, `"cohere"`, `"local"`.
pub fn openai_compat_url(provider: &str) -> Option<&'static str> {
    match provider {
        "openai"  => Some("https://api.openai.com/v1/chat/completions"),
        "gemini"  => Some("https://generativelanguage.googleapis.com/v1beta/openai/chat/completions"),
        "mistral" => Some("https://api.mistral.ai/v1/chat/completions"),
        "groq"    => Some("https://api.groq.com/openai/v1/chat/completions"),
        "cohere"  => Some("https://api.cohere.com/v2/chat"),
        "local"   => Some("http://localhost:11434/v1/chat/completions"),
        _         => None,
    }
}

// ---------------------------------------------------------------------------
// Built-in Anthropic backend
// ---------------------------------------------------------------------------

const ANTHROPIC_API_URL: &str = "https://api.anthropic.com/v1/messages";
const ANTHROPIC_VERSION: &str = "2023-06-01";
const MAX_RETRIES: u32 = 3;

/// LLM completion configuration.
#[derive(Debug, Clone)]
pub struct LlmConfig {
    pub model: String,
    pub max_tokens: u32,
}

impl Default for LlmConfig {
    fn default() -> Self {
        Self {
            model: "claude-haiku-4-5-20251001".to_string(),
            max_tokens: 16_384,
        }
    }
}

/// Anthropic Claude API backend.
///
/// # Example
/// ```no_run
/// use sds_converter_core::converter::llm::{AnthropicBackend, LlmConfig};
/// let backend = AnthropicBackend::new("sk-ant-...", LlmConfig::default());
/// ```
pub struct AnthropicBackend {
    client: Client,
    api_key: String,
    config: LlmConfig,
}

impl AnthropicBackend {
    pub fn new(api_key: impl Into<String>, config: LlmConfig) -> Self {
        Self {
            client: Client::new(),
            api_key: api_key.into(),
            config,
        }
    }
}

impl LlmBackend for AnthropicBackend {
    async fn complete(&self, system: &str, user: &str) -> Result<String, SdsError> {
        // Structured system array enables prompt caching (cache_control: ephemeral).
        // Assistant prefill forces the model to start the JSON object directly.
        // temperature=0 eliminates stochastic section omissions.
        let body = serde_json::json!({
            "model": self.config.model,
            "max_tokens": self.config.max_tokens,
            "temperature": 0,
            "system": [{
                "type": "text",
                "text": system,
                "cache_control": { "type": "ephemeral" }
            }],
            "messages": [
                {"role": "user", "content": user},
                {"role": "assistant", "content": "{"}
            ]
        });

        let mut attempt = 0u32;
        let response = loop {
            let r = self.client
                .post(ANTHROPIC_API_URL)
                .header("x-api-key", &self.api_key)
                .header("anthropic-version", ANTHROPIC_VERSION)
                .header("anthropic-beta", "extended-cache-ttl-2025-04-11")
                .header("content-type", "application/json")
                .json(&body)
                .send()
                .await?;

            let status = r.status().as_u16();
            if (status == 429 || status == 529) && attempt < MAX_RETRIES {
                attempt += 1;
                let secs = 2_u64.pow(attempt);
                tracing::warn!("HTTP {status} (attempt {attempt}/{MAX_RETRIES}), retrying in {secs}s");
                tokio::time::sleep(Duration::from_secs(secs)).await;
            } else {
                break r;
            }
        };

        if !response.status().is_success() {
            let status = response.status().as_u16();
            let message = response.text().await.unwrap_or_default();
            return Err(SdsError::LlmApi { status, message });
        }

        let resp: Value = response.json().await?;
        let text = resp["content"][0]["text"]
            .as_str()
            .ok_or_else(|| SdsError::LlmParse("missing content[0].text".to_string()))?;

        // Prepend the prefill character to reconstruct the complete JSON object.
        Ok(format!("{{{text}"))
    }
}

// ---------------------------------------------------------------------------
// OpenAI-compatible backend (works with OpenAI GPT, Google Gemini, etc.)
// ---------------------------------------------------------------------------

/// Backend for any OpenAI-compatible chat completions API.
///
/// # Example — OpenAI GPT
/// ```no_run
/// use sds_converter_core::converter::llm::{OpenAiCompatBackend, LlmConfig};
/// let config = LlmConfig { model: "gpt-4o".into(), max_tokens: 16384 };
/// let backend = OpenAiCompatBackend::openai("sk-...", config);
/// ```
///
/// # Example — Google Gemini
/// ```no_run
/// use sds_converter_core::converter::llm::{OpenAiCompatBackend, LlmConfig};
/// let config = LlmConfig { model: "gemini-2.0-flash".into(), max_tokens: 16384 };
/// let backend = OpenAiCompatBackend::gemini("AIza...", config);
/// ```
pub struct OpenAiCompatBackend {
    client: Client,
    api_key: String,
    config: LlmConfig,
    base_url: String,
}

impl OpenAiCompatBackend {
    pub fn new(api_key: impl Into<String>, config: LlmConfig, base_url: impl Into<String>) -> Self {
        Self {
            client: Client::new(),
            api_key: api_key.into(),
            config,
            base_url: base_url.into(),
        }
    }

    /// OpenAI GPT backend (api.openai.com).
    pub fn openai(api_key: impl Into<String>, config: LlmConfig) -> Self {
        Self::new(api_key, config, openai_compat_url("openai").unwrap())
    }

    /// Google Gemini backend via OpenAI-compatible endpoint.
    pub fn gemini(api_key: impl Into<String>, config: LlmConfig) -> Self {
        Self::new(api_key, config, openai_compat_url("gemini").unwrap())
    }
}

impl LlmBackend for OpenAiCompatBackend {
    async fn complete(&self, system: &str, user: &str) -> Result<String, SdsError> {
        let body = serde_json::json!({
            "model": self.config.model,
            "max_tokens": self.config.max_tokens,
            "temperature": 0,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user}
            ]
        });

        let mut attempt = 0u32;
        let response = loop {
            let r = self.client
                .post(&self.base_url)
                .header("Authorization", format!("Bearer {}", self.api_key))
                .header("content-type", "application/json")
                .json(&body)
                .send()
                .await?;

            let status = r.status().as_u16();
            if (status == 429 || status == 529) && attempt < MAX_RETRIES {
                attempt += 1;
                let secs = 2_u64.pow(attempt);
                tracing::warn!("HTTP {status} (attempt {attempt}/{MAX_RETRIES}), retrying in {secs}s");
                tokio::time::sleep(Duration::from_secs(secs)).await;
            } else {
                break r;
            }
        };

        if !response.status().is_success() {
            let status = response.status().as_u16();
            let message = response.text().await.unwrap_or_default();
            return Err(SdsError::LlmApi { status, message });
        }

        let resp: Value = response.json().await?;
        let text = resp["choices"][0]["message"]["content"]
            .as_str()
            .ok_or_else(|| SdsError::LlmParse("missing choices[0].message.content".to_string()))?;

        Ok(strip_code_fences(text))
    }
}

// ---------------------------------------------------------------------------
// Shared helpers
// ---------------------------------------------------------------------------

fn strip_code_fences(text: &str) -> String {
    let text = text.trim();
    let text = if let Some(rest) = text.strip_prefix("```json") {
        rest.trim_start_matches(['\n', '\r', ' '])
    } else if let Some(rest) = text.strip_prefix("```") {
        rest.trim_start_matches(['\n', '\r', ' '])
    } else {
        text
    };
    text.strip_suffix("```").unwrap_or(text).trim().to_string()
}

/// Attempt lightweight repair of truncated or malformed JSON before parsing.
///
/// Handles:
/// - Trailing commas before `}` or `]` (common in LLM output)
/// - Unclosed braces/brackets due to context-limit truncation
fn repair_json(s: &str) -> String {
    let mut s = s.to_string();
    loop {
        let rep = s
            .replace(",}", "}")
            .replace(",]", "]")
            .replace(", }", "}")
            .replace(", ]", "]")
            .replace(",\n}", "}")
            .replace(",\n]", "]")
            .replace(",\r\n}", "}")
            .replace(",\r\n]", "]");
        if rep == s {
            break;
        }
        s = rep;
    }

    // Close unclosed braces/brackets using a stack
    let mut stack: Vec<char> = Vec::new();
    let mut in_string = false;
    let mut prev_backslash = false;
    for c in s.chars() {
        if prev_backslash {
            prev_backslash = false;
            continue;
        }
        if c == '\\' && in_string {
            prev_backslash = true;
            continue;
        }
        if c == '"' {
            in_string = !in_string;
            continue;
        }
        if in_string {
            continue;
        }
        match c {
            '{' => stack.push('}'),
            '[' => stack.push(']'),
            '}' | ']' => {
                stack.pop();
            }
            _ => {}
        }
    }
    for closer in stack.iter().rev() {
        s.push(*closer);
    }
    s
}

// ---------------------------------------------------------------------------
// SDS extraction
// ---------------------------------------------------------------------------

const MHLW_SCHEMA_HINT: &str = r#"Output a JSON object. CRITICAL: Use EXACTLY these key names — they must match the MHLW schema precisely.
{
  "Datasheet": { "IssueDate": "YYYY-MM-DD", "SDS-SchemaVersionNo": "1.0" },
  "Identification": {
    "TradeProductIdentity": { "TradeNameJP": "...", "TradeNameEN": "..." },
    "SupplierInformation": { "CompanyName": "...", "Phone": "...", "Address": "..." },
    "RecommendedUseAndRestrictions": "..."
  },
  "HazardIdentification": {
    "Classification": { "FlammableLiquids": "区分2" },
    "HazardLabelling": {
      "SignalWord": "危険",
      "HazardStatement": [{ "HazardStatementCode": "H225", "FullText": "引火性の高い液体および蒸気" }],
      "PrecautionaryStatements": {
        "Prevention": [{ "PrecautionaryStatementCode": "P210", "FullText": "熱から遠ざけること。" }],
        "Response": [{ "PrecautionaryStatementCode": "P370+P378", "FullText": "火災の場合：消火に適切な手段を用いること。" }],
        "Storage": [{ "PrecautionaryStatementCode": "P403+P235", "FullText": "換気の良い場所で保管すること。涼しい場所に置くこと。" }],
        "Disposal": [{ "PrecautionaryStatementCode": "P501", "FullText": "内容物・容器を法規に従って廃棄すること。" }]
      }
    }
  },
  "Composition": {
    "CompositionType": "単一物質",
    "CompositionAndConcentration": [{
      "SubstanceIdentifiers": {
        "SubstanceNames": { "IupacName": "エタノール", "GenericName": "エチルアルコール" },
        "SubstanceIdentity": { "CASno": { "FullText": ["64-17-5"] } }
      },
      "Concentration": { "NumericRangeWithUnitAndQualifier": { "ExactValue": { "Value": 99.5 }, "Unit": "%" } }
    }]
  },
  "FirstAidMeasures": {
    "ExposureRoute": {
      "FirstAidInhalation": { "FullText": "新鮮な空気の場所に移動させ安静を保つ。" },
      "FirstAidSkin": { "FullText": "石鹸と水で洗い流す。" },
      "FirstAidEye": { "FullText": "水で十分に洗い流す。" },
      "FirstAidIngestion": { "FullText": "口をすすぐ。多量に飲み込んだ場合は医師に連絡する。" }
    }
  },
  "FireFightingMeasures": {
    "MediaToBeUsed": { "FullText": "粉末消火剤、二酸化炭素、泡、水噴霧" },
    "FireAndExplosionHazards": { "FullText": "..." },
    "FireFightingProcedures": { "FullText": "..." },
    "SpecialProtectiveEquipmentForFirefighters": { "FullText": "..." }
  },
  "AccidentalReleaseMeasures": {
    "HumanExposureAndEmergencyMeasuress": { "FullText": "..." },
    "EnvironmentalPrecautions": { "FullText": "..." },
    "ContainmentAndCleaningUp": { "FullText": "..." }
  },
  "HandlingAndStorage": {
    "SafeHandling": {
      "HandlingPrecautions": "火気厳禁。換気の良い場所で取り扱う。",
      "TechnicalMeasuresAndStorageConditions": {
        "ProtectiveMeasures": "適切な保護措置を講じる。",
        "VentilationCondition": "局所排気換気を確保する。"
      }
    },
    "Storage": {
      "ConditionsForSafeStorage": { "TechnicalMeasuresAndStorageConditions": "冷暗所に保管する。容器を密閉する。" }
    }
  },
  "ExposureControlPersonalProtection": {
    "AppropriateEngineeringControls": ["局所排気装置を設置する。"],
    "PersonalProtectionEquipment": {
      "EyeProtection": [{ "FullText": "保護眼鏡を着用する。" }],
      "SkinProtection": [{ "FullText": "保護手袋を着用する。" }],
      "RespiratoryProtection": [{ "FullText": "蒸気が高濃度の場合は有機ガス用防毒マスクを着用する。" }],
      "HandProtection": [{ "FullText": "保護手袋を着用する。" }]
    }
  },
  "PhysicalChemicalProperties": {
    "BasePhysicalChemicalProperties": { "PhysicalState": "液体", "Colour": "無色", "Odour": "特異臭" },
    "MeltingPointRelated": [{ "ItemName": "融点", "NumericRangeWithUnitAndQualifier": { "ExactValue": { "Value": -117.0 }, "Unit": "°C" } }],
    "BoilingPointRelated": [{ "ItemName": "沸点", "NumericRangeWithUnitAndQualifier": { "ExactValue": { "Value": 78.4 }, "Unit": "°C" } }],
    "FlashPoint": [{ "NumericRangeWithUnitAndQualifier": { "ExactValue": { "Value": 13.0 }, "Unit": "°C" } }],
    "VapourPressure": [{ "NumericRangeWithUnitAndQualifier": { "ExactValue": { "Value": 5.9 }, "Unit": "kPa" } }],
    "Densities": [{ "NumericRangeWithUnitAndQualifier": { "ExactValue": { "Value": 0.789 }, "Unit": "g/cm3" } }],
    "Solubilities": { "WaterSolubility": [{ "AdditionalInfo": { "FullText": ["水に混和する"] } }] },
    "OtherPhysicalChemicalProperty": [{ "ItemName": "その他の物性", "AdditionalInfo": { "FullText": ["テキスト形式の補足情報"] } }]
  },
  "StabilityReactivity": {
    "ReactivityDescription": "反応性についての説明",
    "StabilityDescription": "通常の保管条件で安定",
    "HazardousReactions": { "FullText": "強酸化剤と反応する。" },
    "ConditionsToAvoid": "熱、火花、炎",
    "MaterialsToAvoid": "強酸化剤",
    "HazardousDecompositionProducts": { "Substance": "炭素酸化物" }
  },
  "ToxicologicalInformation": [{
    "AcuteToxicity": { "ExposureRoute": [{ "RouteOfExposure": "経口", "FullText": "LD50 ラット 経口 7060 mg/kg" }] },
    "SkinCorrosionIrritation": { "TestResult": [{ "FullText": "軽度の皮膚刺激性" }] },
    "EyeDamageOrIrritation": { "TestResult": [{ "FullText": "軽度の眼刺激性" }] },
    "AdditionalToxicologicalInformation": "その他の毒性情報"
  }],
  "EcologicalInformation": [{
    "EcotoxicologicalInformation": { "AquaticToxicity": [{ "FullText": "..." }] },
    "PersistenceDegradability": { "FullText": "生分解性あり" },
    "AdditionalEcotoxInformation": "その他の生態影響情報"
  }],
  "DisposalConsiderations": {
    "FullText": "廃棄物は関係法令に従って処理する。",
    "ProductWaste": ["廃棄物処理法に従い処理する"],
    "PackagingWaste": ["容器はリサイクルまたは適切に廃棄する"]
  },
  "TransportInformation": {
    "InternationalRegulations": [{
      "RegulationName": [{ "TransportationType": "海上輸送(IMDG)", "FullText": "国連番号：UN1170 品名：ETHANOL" }]
    }],
    "DomesticRegulations": [{ "TransportationType": "陸上輸送", "FullText": "消防法：第4類 アルコール類" }]
  },
  "RegulatoryInformation": {
    "OtherLegislation": {
      "Legislation": [{ "LegislationName": "労働安全衛生法", "Regulations": [{ "FullText": "..." }] }]
    }
  },
  "OtherInformation": {
    "RelatedDocuments": "参考資料リスト",
    "RevisionInformation": [{ "LastUpdateDate": "YYYY-MM-DD", "FullText": "改訂内容の説明" }],
    "Desclaimer": "免責事項テキスト"
  }
}"#;

const TYPO_WARNINGS: &str = r#"
CRITICAL — the following key names look like typos but are REQUIRED exactly as shown in the MHLW schema:
- "HumanExposureAndEmergencyMeasuress"  (two s's at the end — do NOT change to single s)
- "TestGuidline"  (not "TestGuideline" — the 'e' is missing intentionally)
- "Desclaimer"  (not "Disclaimer" — letters are transposed intentionally)
- "gazetteNo"  (lowercase 'g' — do NOT capitalize)
- "Dose/Concentration"  (literal forward-slash in the key name)
Use these malformed spellings exactly. Correcting them will break schema validation."#;

/// Build the system prompt for SDS extraction in the given document language.
fn build_system_prompt(lang: Option<Language>) -> String {
    let lang_hint = match lang {
        Some(l) => format!(
            "The source document is written in {} ({}).\n",
            l.name_en(),
            l.bcp47()
        ),
        None => "The source document may be in Japanese, English, Simplified Chinese, or Traditional Chinese — detect the language automatically.\n".to_string(),
    };

    let section_hint = match lang {
        Some(Language::Japanese) | None => {
            "Section headings follow JIS Z 7253 (第1節〜第16節).\n"
        }
        Some(Language::English) => {
            "Section headings follow GHS/OSHA HazCom (SECTION 1–16).\n"
        }
        Some(Language::ChineseSimplified) => {
            "Section headings follow GB/T 16483 (第1部分〜第16部分).\n"
        }
        Some(Language::ChineseTraditional) => {
            "Section headings follow CNS 15030 (第1節〜第16節).\n"
        }
    };

    format!(
        "You are an expert in extracting Safety Data Sheet (SDS) information.\n\
         {lang_hint}\
         {section_hint}\
         Read the document text and output all SDS information as a JSON object conforming to the \
         Japanese Ministry of Health, Labour and Welfare (MHLW) SDS data exchange format v1.0.\n\
         Rules:\n\
         - Output raw JSON only — no markdown, no code fences, no explanation\n\
         - Your response must begin immediately with '{{' — the first character must be '{{'\n\
         - CRITICAL: Extract ALL sections listed in the user message. Never silently omit a section.\n\
         - Pay special attention to Section 9 (PhysicalChemicalProperties): always include it if the document has any physical/chemical property data, even if only BasePhysicalChemicalProperties\n\
         - For Section 9 numeric properties (FlashPoint, VapourPressure, Densities, etc.): use NumericRangeWithUnitAndQualifier with a numeric Value. If the value is text only (e.g. '不明', 'N/A', 'データなし'), use AdditionalInfo: {{\"FullText\": [\"text\"]}} instead — never put text in a numeric Value field\n\
         - Omit keys that have no information (empty strings, null, and empty objects {{}} are forbidden)\n\
         - Dates in YYYY-MM-DD format\n\
         - Numeric values as numeric types (not strings) inside NumericRangeWithUnitAndQualifier\n\
         - For qualitative text values in PhysicalChemicalProperties, use AdditionalInfo: {{\"FullText\": [\"text\"]}} — note FullText is an ARRAY of strings\n\
         - Reproduce text exactly as written in the source document; do not infer or fill in missing data\n\
         - JSON keys must match EXACTLY the key names shown in the schema example below\n\
         {TYPO_WARNINGS}\n\
         \nSchema example (use these EXACT key names):\n{MHLW_SCHEMA_HINT}"
    )
}

// Section groups for parallel extraction — splits output token load across two concurrent calls.
const GROUP_A: &[&str] = &[
    "Datasheet",
    "Identification",
    "HazardIdentification",
    "Composition",
    "FirstAidMeasures",
    "FireFightingMeasures",
    "AccidentalReleaseMeasures",
    "HandlingAndStorage",
    "ExposureControlPersonalProtection",
];
const GROUP_B: &[&str] = &[
    "PhysicalChemicalProperties",
    "StabilityReactivity",
    "ToxicologicalInformation",
    "EcologicalInformation",
    "DisposalConsiderations",
    "TransportInformation",
    "RegulatoryInformation",
    "OtherInformation",
];

/// Merge two `SdsRoot` values by taking the first non-`None` for each field.
fn merge_sds(a: SdsRoot, b: SdsRoot) -> SdsRoot {
    SdsRoot {
        datasheet: a.datasheet.or(b.datasheet),
        identification: a.identification.or(b.identification),
        hazard_identification: a.hazard_identification.or(b.hazard_identification),
        composition: a.composition.or(b.composition),
        first_aid_measures: a.first_aid_measures.or(b.first_aid_measures),
        fire_fighting_measures: a.fire_fighting_measures.or(b.fire_fighting_measures),
        accidental_release_measures: a.accidental_release_measures.or(b.accidental_release_measures),
        handling_and_storage: a.handling_and_storage.or(b.handling_and_storage),
        exposure_control_personal_protection: a
            .exposure_control_personal_protection
            .or(b.exposure_control_personal_protection),
        physical_chemical_properties: a.physical_chemical_properties.or(b.physical_chemical_properties),
        stability_reactivity: a.stability_reactivity.or(b.stability_reactivity),
        toxicological_information: a.toxicological_information.or(b.toxicological_information),
        ecological_information: a.ecological_information.or(b.ecological_information),
        disposal_considerations: a.disposal_considerations.or(b.disposal_considerations),
        transport_information: a.transport_information.or(b.transport_information),
        regulatory_information: a.regulatory_information.or(b.regulatory_information),
        other_information: a.other_information.or(b.other_information),
    }
}

/// Extract SDS data from document text using the provided LLM backend.
///
/// Issues two parallel LLM calls (sections 1–9 and 10–16) to halve per-file latency,
/// then retries any sections skipped due to schema mismatch.
///
/// Returns `(SdsRoot, Vec<String>)` where the `Vec` lists any sections that could not
/// be extracted after all passes.
pub async fn extract_sds_from_text<B: LlmBackend + Sync>(
    backend: &B,
    text: &str,
    source_language: Option<Language>,
) -> Result<(SdsRoot, Vec<String>), SdsError> {
    let system = build_system_prompt(source_language);

    let lang_prefix = match source_language {
        Some(l) => format!("This document is in {}. ", l.name_en()),
        None => String::new(),
    };

    let user_a = format!(
        "{lang_prefix}Extract ONLY these sections: {}.\n\
         Output as JSON. Do not include any other sections.\n\n\
         Document text:\n{text}",
        GROUP_A.join(", ")
    );
    let user_b = format!(
        "{lang_prefix}Extract ONLY these sections: {}.\n\
         Output as JSON. Do not include any other sections.\n\n\
         Document text:\n{text}",
        GROUP_B.join(", ")
    );

    // Parallel extraction of both groups — each call generates ~half the output tokens.
    let (raw_a, raw_b) = tokio::join!(
        backend.complete(&system, &user_a),
        backend.complete(&system, &user_b),
    );

    let json_a = raw_a?;
    let json_b = raw_b?;
    tracing::debug!("Group A JSON:\n{json_a}");
    tracing::debug!("Group B JSON:\n{json_b}");

    let (sds_a, skipped_a) = lenient_deserialize(&json_a)?;
    let (sds_b, skipped_b) = lenient_deserialize(&json_b)?;
    let mut sds = merge_sds(sds_a, sds_b);
    let mut all_skipped = [skipped_a, skipped_b].concat();

    // Retry pass: re-request only the sections that had schema issues.
    if !all_skipped.is_empty() {
        let retry_keys: Vec<&str> = all_skipped.iter().map(String::as_str).collect();
        tracing::warn!(
            "Retrying {} skipped sections: {}",
            retry_keys.len(),
            retry_keys.join(", ")
        );
        let user_retry = format!(
            "{lang_prefix}Extract ONLY these sections (previous extraction had schema issues — \
             be especially precise about field types and nesting): {}.\n\
             Output as JSON.\n\n\
             Document text:\n{text}",
            retry_keys.join(", ")
        );
        if let Ok(raw_retry) = backend.complete(&system, &user_retry).await {
            tracing::debug!("Retry JSON:\n{raw_retry}");
            if let Ok((retry_sds, retry_skipped)) = lenient_deserialize(&raw_retry) {
                sds = merge_sds(sds, retry_sds);
                all_skipped = retry_skipped;
            }
        }
    }

    let warnings: Vec<String> = all_skipped
        .into_iter()
        .map(|k| format!("{k}: skipped (schema mismatch — check logs for details)"))
        .collect();

    Ok((sds, warnings))
}

/// Deserialize LLM JSON output section-by-section, skipping sections with type errors.
///
/// Returns `(SdsRoot, Vec<String>)` where the Vec lists skipped section *key names*
/// (e.g. `["HandlingAndStorage"]`). The full serde error is written to the tracing log.
fn lenient_deserialize(json_str: &str) -> Result<(SdsRoot, Vec<String>), SdsError> {
    use crate::schema::*;
    use tracing::warn;

    let repaired = repair_json(json_str);

    let val: Value = serde_json::from_str(&repaired).map_err(|e| {
        let preview: String = json_str.chars().take(500).collect();
        SdsError::LlmParse(format!("Invalid JSON: {e}\nRaw (first 500 chars): {preview}"))
    })?;

    let obj = val
        .as_object()
        .ok_or_else(|| SdsError::LlmParse("LLM output is not a JSON object".into()))?;

    let mut skipped: Vec<String> = Vec::new();

    macro_rules! section {
        ($key:literal, $type:ty) => {
            obj.get($key).and_then(|v| {
                serde_json::from_value::<$type>(v.clone())
                    .map_err(|e| {
                        let preview: String = serde_json::to_string(v)
                            .unwrap_or_default()
                            .chars()
                            .take(200)
                            .collect();
                        warn!(
                            "Section '{}' skipped (schema mismatch): {}\nValue preview: {}",
                            $key, e, preview
                        );
                        skipped.push($key.to_string());
                    })
                    .ok()
            })
        };
    }

    let sds = SdsRoot {
        datasheet: section!("Datasheet", Datasheet),
        identification: section!("Identification", Identification),
        hazard_identification: section!("HazardIdentification", HazardIdentification),
        composition: section!("Composition", Composition),
        first_aid_measures: section!("FirstAidMeasures", FirstAidMeasures),
        fire_fighting_measures: section!("FireFightingMeasures", FireFightingMeasures),
        accidental_release_measures: section!("AccidentalReleaseMeasures", AccidentalReleaseMeasures),
        handling_and_storage: section!("HandlingAndStorage", HandlingAndStorage),
        exposure_control_personal_protection: section!(
            "ExposureControlPersonalProtection",
            ExposureControlPersonalProtection
        ),
        physical_chemical_properties: section!("PhysicalChemicalProperties", PhysicalChemicalProperties),
        stability_reactivity: section!("StabilityReactivity", StabilityReactivity),
        toxicological_information: section!("ToxicologicalInformation", Vec<ToxicologicalInformation>),
        ecological_information: section!("EcologicalInformation", Vec<EcologicalInformation>),
        disposal_considerations: section!("DisposalConsiderations", DisposalConsiderations),
        transport_information: section!("TransportInformation", TransportInformation),
        regulatory_information: section!("RegulatoryInformation", RegulatoryInformation),
        other_information: section!("OtherInformation", OtherInformation),
    };

    Ok((sds, skipped))
}

// ---------------------------------------------------------------------------
// Unit tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn strip_fences_bare_json() {
        assert_eq!(strip_code_fences("{}"), "{}");
    }

    #[test]
    fn strip_fences_json_tag() {
        assert_eq!(strip_code_fences("```json\n{}\n```"), "{}");
    }

    #[test]
    fn strip_fences_plain_backticks() {
        assert_eq!(strip_code_fences("```\n{}\n```"), "{}");
    }

    #[test]
    fn strip_fences_no_fences() {
        let raw = r#"{"key": "value"}"#;
        assert_eq!(strip_code_fences(raw), raw);
    }

    #[test]
    fn strip_fences_whitespace_after_tag() {
        assert_eq!(strip_code_fences("```json  \n{}\n```"), "{}");
    }

    #[test]
    fn repair_trailing_comma_before_brace() {
        let input = r#"{"a": 1,}"#;
        let result = repair_json(input);
        serde_json::from_str::<serde_json::Value>(&result).expect("should be valid JSON");
    }

    #[test]
    fn repair_trailing_comma_before_bracket() {
        let input = r#"{"a": [1, 2,]}"#;
        let result = repair_json(input);
        serde_json::from_str::<serde_json::Value>(&result).expect("should be valid JSON");
    }

    #[test]
    fn repair_unclosed_object() {
        let input = r#"{"a": 1, "b": {"c": 2}"#;
        let result = repair_json(input);
        serde_json::from_str::<serde_json::Value>(&result).expect("should be valid JSON");
    }

    #[test]
    fn repair_already_valid_unchanged() {
        let input = r#"{"a": 1}"#;
        let result = repair_json(input);
        assert_eq!(result, input);
    }

    #[test]
    fn openai_compat_url_known_providers() {
        assert!(openai_compat_url("openai").is_some());
        assert!(openai_compat_url("gemini").is_some());
        assert!(openai_compat_url("mistral").is_some());
        assert!(openai_compat_url("groq").is_some());
        assert!(openai_compat_url("cohere").is_some());
        assert!(openai_compat_url("local").is_some());
        assert!(openai_compat_url("unknown").is_none());
    }
}