multiagent_template_extractor.py

"""
Multi-Agent Template-Based Data Extractor
Combines template-driven field definitions with multi-agent extraction pipeline
for accurate extraction from spec sheets with multiple variants/models
"""

import json
import fitz  # PyMuPDF
from pathlib import Path
from typing import Dict, List, Optional, Any

# Import unified LLM client
from llm_client import chat_completion, validate_env, get_active_provider


def extract_text_from_pdf(pdf_path: str) -> Optional[str]:
    """Extract all text from a PDF file"""
    try:
        pdf_document = fitz.open(pdf_path)
        extracted_text = ""

        for page_num in range(pdf_document.page_count):
            page = pdf_document.load_page(page_num)
            text = page.get_text()
            extracted_text += f"\n--- Page {page_num + 1} ---\n{text}"

        pdf_document.close()
        return extracted_text

    except Exception as e:
        print(f"Error extracting text from PDF: {e}")
        return None


def parse_json_response(response: str) -> Optional[Any]:
    """Parse JSON from LLM response, handling markdown code blocks and extra text"""
    if response is None:
        return None

    # Try to find JSON block in markdown code fence
    if '```json' in response:
        start = response.find('```json') + 7
        end = response.find('```', start)
        if end > start:
            response = response[start:end].strip()
    elif '```' in response:
        start = response.find('```') + 3
        end = response.find('```', start)
        if end > start:
            response = response[start:end].strip()

    # If no code block, try to find JSON object or array directly
    if not response.startswith('{') and not response.startswith('['):
        start_obj = response.find('{')
        start_arr = response.find('[')

        if start_obj == -1:
            start = start_arr
        elif start_arr == -1:
            start = start_obj
        else:
            start = min(start_obj, start_arr)

        end_obj = response.rfind('}')
        end_arr = response.rfind(']')
        end = max(end_obj, end_arr)

        if start != -1 and end > start:
            response = response[start:end+1]

    response = response.strip()

    try:
        return json.loads(response)
    except json.JSONDecodeError as e:
        print(f"JSON parsing error: {e}")
        print(f"Response text: {response[:500]}")
        return None


class MultiAgentTemplateExtractor:
    """
    Multi-agent extractor that uses JSON templates to define extraction schemas
    and a multi-step pipeline for accurate extraction from multi-variant spec sheets
    """

    def __init__(self, template_path: str):
        """
        Initialize the extractor with a template file

        Args:
            template_path: Path to the JSON template file
        """
        self.template = self._load_template(template_path)
        self.component_type = self.template.get('component_type', 'unknown')

    def _load_template(self, template_path: str) -> Dict:
        """Load and validate the template file"""
        try:
            with open(template_path, 'r') as f:
                template = json.load(f)

            # Validate required template fields
            required_keys = ['component_type', 'fields']
            for key in required_keys:
                if key not in template:
                    raise ValueError(f"Template missing required key: {key}")

            return template

        except FileNotFoundError:
            raise FileNotFoundError(f"Template file not found: {template_path}")
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in template file: {e}")

    # =========================================================================
    # AGENT 1: DOCUMENT ANALYZER
    # =========================================================================

    def _analyze_document_structure(self, pdf_text: str, llm_model: str) -> Dict:
        """
        Agent 1: Analyze document to identify structure and variants
        """
        variant_config = self.template.get('variant_detection', {})
        variant_enabled = variant_config.get('enabled', False)
        variant_key = variant_config.get('variant_key', 'model')
        variant_desc = variant_config.get('variant_description', 'models')
        variant_hint = variant_config.get('prompt_hint', '')
        variant_unit = variant_config.get('variant_unit', '')

        component_type = self.template.get('component_type', 'component')
        description = self.template.get('description', '')

        prompt = f"""Analyze this {component_type} specification sheet and identify its structure.

{description}

PDF Text (first 8000 characters):
{pdf_text[:8000]}

Determine:
1. The manufacturer name
2. The model series or product line name
3. Whether this document contains MULTIPLE {variant_desc} or just ONE
4. If multiple, list all the {variant_desc} found

{variant_hint}

Return ONLY a JSON response with this exact structure:
{{
  "manufacturer": "manufacturer name",
  "model_series": "model series or product line name",
  "document_type": "single_variant" or "multi_variant",
  "variants": [list of {variant_key} values found, e.g., {f'[350, 355, 360, 365, 370, 375]' if variant_unit else '["Model-A", "Model-B", "Model-C"]'}],
  "num_variants": integer count of variants found
}}

Be precise. Look carefully for tables or columns that list multiple {variant_desc}.
Return ONLY the JSON, no explanations."""

        result = chat_completion(
            messages=[{"role": "user", "content": prompt}],
            model=llm_model,
            max_tokens=1000,
            temperature=0.0
        )

        structure = parse_json_response(result)

        if structure is None:
            structure = {
                "manufacturer": "Unknown",
                "model_series": "Unknown",
                "document_type": "unknown",
                "variants": [],
                "num_variants": 0
            }

        return structure

    # =========================================================================
    # AGENT 2: SHARED PARAMETER EXTRACTOR
    # =========================================================================

    def _extract_shared_parameters(self, pdf_text: str, llm_model: str) -> Dict:
        """
        Agent 2: Extract parameters that are shared across all variants
        (e.g., temperature coefficient units for PV modules)
        """
        shared_config = self.template.get('shared_parameters', {})

        if not shared_config.get('enabled', False):
            return {}

        shared_fields = shared_config.get('fields', [])
        if not shared_fields:
            return {}

        description = shared_config.get('description', 'Shared parameters')

        # Build field extraction prompt
        field_prompts = []
        for field in shared_fields:
            name = field['name']
            desc = field['description']
            hint = field.get('prompt_hint', '')
            allowed = field.get('allowed_values', [])

            field_str = f"- {name}: {desc}"
            if allowed:
                field_str += f" (choose from: {allowed})"
            if hint:
                field_str += f"\n  Hint: {hint}"
            field_prompts.append(field_str)

        prompt = f"""Extract the following SHARED parameters from this specification sheet.

{description}

These values are the SAME for all variants in the document - they are typically shown once.

PDF Text:
{pdf_text[:8000]}

Extract these parameters:
{chr(10).join(field_prompts)}

Return ONLY a JSON object with these exact keys.
If a value is not found, use "N/A".

Example format:
{{
  "{shared_fields[0]['name']}": "value",
  ...
}}"""

        result = chat_completion(
            messages=[{"role": "user", "content": prompt}],
            model=llm_model,
            max_tokens=500,
            temperature=0.0
        )

        shared_params = parse_json_response(result)
        return shared_params if shared_params else {}

    # =========================================================================
    # AGENT 3: VARIANT EXTRACTOR
    # =========================================================================

    def _extract_single_variant(self, pdf_text: str, variant_id: Any,
                                 structure: Dict, shared_params: Dict,
                                 llm_model: str) -> Optional[Dict]:
        """
        Agent 3: Extract parameters for a specific variant
        """
        variant_config = self.template.get('variant_detection', {})
        variant_key = variant_config.get('variant_key', 'model')
        variant_unit = variant_config.get('variant_unit', '')

        fields = self.template.get('fields', [])
        test_conditions = self.template.get('test_conditions')
        example = self.template.get('example_output', {})

        # Group fields by category
        categories = {}
        for field in fields:
            category = field.get('category', 'parameters')
            if category not in categories:
                categories[category] = []
            categories[category].append(field)

        # Build field descriptions by category
        field_sections = []
        for category, cat_fields in categories.items():
            section = f"\n{category.upper().replace('_', ' ')}:"
            for field in cat_fields:
                name = field['name']
                desc = field['description']
                required = field.get('required', False)
                default = field.get('default', 'N/A')
                allowed = field.get('allowed_values')

                # Skip shared parameter unit fields - they come from Agent 2
                if name in shared_params:
                    continue

                field_str = f"\n- {name}: {desc}"
                if allowed:
                    field_str += f" (choose from: {allowed})"
                if not required:
                    field_str += f' (use "{default}" if not found)'

                section += field_str
            field_sections.append(section)

        variant_display = f"{variant_id}{variant_unit}" if variant_unit else variant_id

        prompt = f"""Extract specifications for the {variant_display} variant ONLY from this datasheet.

Manufacturer: {structure.get('manufacturer', 'Unknown')}
Model Series: {structure.get('model_series', 'Unknown')}
Target Variant: {variant_display}

PDF Text:
{pdf_text[:10000]}

Extract these parameters FOR THE {variant_display} VARIANT ONLY:
{''.join(field_sections)}
"""

        if test_conditions:
            prompt += f"\nIMPORTANT: Only include data from {test_conditions}.\n"

        # Add info about shared parameters if any
        if shared_params:
            prompt += f"\nNote: The following shared parameters will be added automatically: {list(shared_params.keys())}\n"

        prompt += """
CRITICAL INSTRUCTIONS:
1. Look in tables for the row/column corresponding to this specific variant
2. If a value is not found, use "N/A"
3. All numerical values must be numbers, not strings (except "N/A")
4. Double-check you're extracting from the correct column/row

Return ONLY a JSON object with the field names as keys."""

        if example:
            # Show a simplified example
            example_subset = {k: v for k, v in list(example.items())[:5]}
            prompt += f"\n\nExample format:\n{json.dumps(example_subset, indent=2)}"

        result = chat_completion(
            messages=[{"role": "user", "content": prompt}],
            model=llm_model,
            max_tokens=2000,
            temperature=0.0
        )

        variant_data = parse_json_response(result)

        if variant_data is None:
            return None

        # Add shared parameters
        variant_data.update(shared_params)

        return variant_data

    # =========================================================================
    # AGENT 4: DATA VALIDATOR
    # =========================================================================

    def _validate_data(self, data_dict: Dict) -> Dict:
        """
        Agent 4: Validate extracted data against template rules
        """
        validation_rules = self.template.get('validation_rules', [])

        validation = {
            "valid": True,
            "errors": [],
            "warnings": []
        }

        for rule in validation_rules:
            rule_name = rule.get('rule', 'unknown')
            description = rule.get('description', '')
            severity = rule.get('severity', 'warning')

            try:
                # Range validation
                if 'range' in rule:
                    field_name = rule.get('field')
                    if field_name and field_name in data_dict and data_dict[field_name] != 'N/A':
                        value = data_dict[field_name]
                        min_val, max_val = rule['range']

                        if not (min_val <= value <= max_val):
                            msg = f"{field_name} ({value}) outside valid range [{min_val}, {max_val}]"
                            if severity == 'error':
                                validation['errors'].append(msg)
                                validation['valid'] = False
                            else:
                                validation['warnings'].append(msg)

                # Condition validation
                elif 'condition' in rule:
                    condition = rule['condition']
                    # Build safe context with only numeric values
                    context = {k: v for k, v in data_dict.items() if v != 'N/A' and isinstance(v, (int, float))}

                    if len(context) > 0:
                        try:
                            result = eval(condition, {"__builtins__": {}, "abs": abs}, context)
                            if not result:
                                msg = f"{description}"
                                if severity == 'error':
                                    validation['errors'].append(msg)
                                    validation['valid'] = False
                                else:
                                    validation['warnings'].append(msg)
                        except:
                            # Skip validation if required fields are missing
                            pass

                # Allowed values validation
                elif 'allowed_values' in rule:
                    field = rule.get('field')
                    if field and field in data_dict and data_dict[field] != 'N/A':
                        if data_dict[field] not in rule['allowed_values']:
                            msg = f"{field} value '{data_dict[field]}' not in allowed values"
                            if severity == 'error':
                                validation['errors'].append(msg)
                                validation['valid'] = False
                            else:
                                validation['warnings'].append(msg)

            except Exception as e:
                print(f"Warning: Could not apply validation rule '{rule_name}': {e}")

        return validation

    # =========================================================================
    # AGENT 5: ERROR CORRECTOR
    # =========================================================================

    def _correct_extraction_errors(self, pdf_text: str, variant_data: Dict,
                                    validation: Dict, variant_id: Any,
                                    llm_model: str) -> Dict:
        """
        Agent 5: Re-extract fields that failed validation
        """
        if not validation['errors']:
            return {}

        # Identify fields that need re-extraction from error messages
        error_fields = set()
        for error in validation['errors']:
            # Try to identify field names in error message
            for field in self.template.get('fields', []):
                if field['name'] in error:
                    error_fields.add(field['name'])

        if not error_fields:
            return {}

        variant_config = self.template.get('variant_detection', {})
        variant_unit = variant_config.get('variant_unit', '')
        variant_display = f"{variant_id}{variant_unit}" if variant_unit else variant_id

        prompt = f"""The following values were extracted but FAILED VALIDATION for {variant_display}:

Validation Errors:
{chr(10).join('- ' + e for e in validation['errors'])}

Current extracted values that need correction:
{json.dumps({k: variant_data.get(k, 'N/A') for k in error_fields}, indent=2)}

Please re-extract ONLY these specific values from the PDF:
{', '.join(error_fields)}

PDF Text:
{pdf_text[:10000]}

CRITICAL:
1. Look VERY CAREFULLY at the table for the {variant_display} column/row
2. Double-check you're reading the correct cell
3. Verify the values make physical sense

Return ONLY a JSON with the corrected values:
{{
  "field_name": corrected_value,
  ...
}}

All values must be numbers (not strings)."""

        result = chat_completion(
            messages=[{"role": "user", "content": prompt}],
            model=llm_model,
            max_tokens=800,
            temperature=0.0
        )

        corrections = parse_json_response(result)
        return corrections if corrections else {}

    # =========================================================================
    # NORMALIZATION
    # =========================================================================

    def _normalize_data(self, data_dict: Dict) -> Dict:
        """
        Apply normalization rules from the template
        """
        normalization_rules = self.template.get('normalization_rules', [])

        if not normalization_rules:
            return data_dict

        for rule in normalization_rules:
            field = rule['field']
            from_unit = rule['from_unit']
            to_unit = rule['to_unit']
            formula = rule['formula']

            # Check if this field exists and needs normalization
            unit_field = f"{field}_unit"

            if (field in data_dict and
                unit_field in data_dict and
                data_dict[field] != 'N/A' and
                data_dict[unit_field] == from_unit):

                # Build context for formula evaluation
                context = {k: v for k, v in data_dict.items() if v != 'N/A' and isinstance(v, (int, float))}
                context['value'] = data_dict[field]

                try:
                    new_value = eval(formula, {"__builtins__": {}}, context)
                    data_dict[field] = round(new_value, 4)
                    data_dict[unit_field] = to_unit
                except Exception as e:
                    print(f"Warning: Could not apply normalization rule for {field}: {e}")

        return data_dict

    # =========================================================================
    # LOGGING HELPER
    # =========================================================================

    def _log(self, message: str, callback=None):
        """Log a message to console and optionally to a callback"""
        print(message)
        if callback:
            callback(message)

    # =========================================================================
    # MAIN EXTRACTION PIPELINE
    # =========================================================================

    def extract(self, pdf_path: str, llm_model: str = "openai/gpt-oss-120b",
                normalize: bool = True, validate: bool = True,
                log_callback=None) -> Dict:
        """
        Multi-agent extraction pipeline

        Args:
            pdf_path: Path to the PDF file
            llm_model: LLM model to use for extraction
            normalize: Whether to apply normalization rules
            validate: Whether to validate extracted data
            log_callback: Optional callback function for log messages

        Returns:
            dict with keys: success, data, structure, validation, error
        """
        # Validate environment
        if not validate_env():
            return {
                'success': False,
                'data': [],
                'structure': None,
                'validation': None,
                'error': 'LLM environment not configured'
            }

        log = lambda msg: self._log(msg, log_callback)

        log(f"\n{'='*70}")
        log(f"MULTI-AGENT {self.component_type.upper()} EXTRACTION")
        log(f"{'='*70}")
        log(f"PDF: {pdf_path}")
        log(f"Template: {self.template.get('description', '')}")
        log(f"LLM Model: {llm_model}")
        log(f"Provider: {get_active_provider() or 'auto-detect'}\n")

        # Extract text from PDF
        pdf_text = extract_text_from_pdf(pdf_path)

        if pdf_text is None or len(pdf_text.strip()) < 100:
            return {
                'success': False,
                'data': [],
                'structure': None,
                'validation': None,
                'error': 'Could not extract sufficient text from PDF'
            }

        log(f"Extracted {len(pdf_text)} characters from PDF\n")

        # =====================================================================
        # STEP 1: ANALYZE DOCUMENT STRUCTURE
        # =====================================================================
        log("STEP 1: Analyzing document structure...")
        log("-" * 70)

        structure = self._analyze_document_structure(pdf_text, llm_model)

        log(f"  Manufacturer: {structure.get('manufacturer', 'Unknown')}")
        log(f"  Model Series: {structure.get('model_series', 'Unknown')}")
        log(f"  Document Type: {structure.get('document_type', 'Unknown')}")
        log(f"  Variants Found: {structure.get('num_variants', 0)}")
        log(f"  Variants: {structure.get('variants', [])}")
        log("")

        variants = structure.get('variants', [])
        if not variants:
            # Fall back to single extraction
            variants = [None]

        # =====================================================================
        # STEP 2: EXTRACT SHARED PARAMETERS
        # =====================================================================
        shared_config = self.template.get('shared_parameters', {})
        shared_params = {}

        if shared_config.get('enabled', False):
            log("STEP 2: Extracting shared parameters...")
            log("-" * 70)

            shared_params = self._extract_shared_parameters(pdf_text, llm_model)

            for key, value in shared_params.items():
                log(f"  {key}: {value}")
            log("")
        else:
            log("STEP 2: No shared parameters defined (skipped)")
            log("")

        # =====================================================================
        # STEP 3: EXTRACT EACH VARIANT
        # =====================================================================
        log("STEP 3: Extracting variants...")
        log("-" * 70)

        extracted_items = []

        for i, variant_id in enumerate(variants):
            display_name = variant_id if variant_id else "Primary"
            log(f"\n  [{i+1}/{len(variants)}] Extracting {display_name}...")

            variant_data = self._extract_single_variant(
                pdf_text, variant_id, structure, shared_params, llm_model
            )

            if variant_data:
                extracted_items.append(variant_data)
                log(f"    ✓ Extraction complete")
            else:
                log(f"    ✗ Extraction failed")

        if not extracted_items:
            return {
                'success': False,
                'data': [],
                'structure': structure,
                'validation': None,
                'error': 'Failed to extract any variants'
            }

        log("")

        # =====================================================================
        # STEP 4: VALIDATE AND CORRECT
        # =====================================================================
        validations = []

        if validate:
            log("STEP 4: Validating and correcting...")
            log("-" * 70)

            for i, item in enumerate(extracted_items):
                variant_id = variants[i] if i < len(variants) else None
                display_name = variant_id if variant_id else "Primary"
                log(f"\n  [{i+1}/{len(extracted_items)}] Validating {display_name}...")

                validation = self._validate_data(item)

                if validation['errors']:
                    log(f"    ⚠ Validation errors:")
                    for error in validation['errors']:
                        log(f"      - {error}")

                    # Attempt correction
                    log(f"    → Attempting error correction...")
                    corrections = self._correct_extraction_errors(
                        pdf_text, item, validation, variant_id, llm_model
                    )

                    if corrections:
                        item.update(corrections)
                        log(f"    ✓ Applied {len(corrections)} correction(s)")

                        # Re-validate
                        validation = self._validate_data(item)
                        if not validation['errors']:
                            log(f"    ✓ Corrections successful")
                        else:
                            log(f"    ⚠ Some errors remain")
                    else:
                        log(f"    ✗ No corrections generated")
                else:
                    log(f"    ✓ Validation passed")

                if validation['warnings']:
                    log(f"    ⚠ Warnings:")
                    for warning in validation['warnings']:
                        log(f"      - {warning}")

                validations.append(validation)

            log("")

        # =====================================================================
        # STEP 5: NORMALIZE
        # =====================================================================
        if normalize:
            log("STEP 5: Normalizing units...")
            log("-" * 70)

            for i, item in enumerate(extracted_items):
                variant_id = variants[i] if i < len(variants) else None
                display_name = variant_id if variant_id else "Primary"
                self._normalize_data(item)
                log(f"  [{i+1}/{len(extracted_items)}] ✓ Normalized {display_name}")

            log("")

        # =====================================================================
        # FINAL SUMMARY
        # =====================================================================
        log("=" * 70)
        log("EXTRACTION COMPLETE")
        log("=" * 70)
        log(f"Successfully extracted {len(extracted_items)} {self.component_type}(s)")
        log("")

        return {
            'success': True,
            'data': extracted_items,
            'structure': structure,
            'validation': validations if validate else None,
            'error': None
        }


# =============================================================================
# CONVENIENCE FUNCTIONS
# =============================================================================

def extract_from_pdf(pdf_path: str, template_path: str,
                     llm_model: str = "openai/gpt-oss-120b",
                     normalize: bool = True, validate: bool = True,
                     log_callback=None) -> Dict:
    """
    Quick extraction function using a template

    Args:
        pdf_path: Path to PDF file
        template_path: Path to template JSON file
        llm_model: LLM model to use
        normalize: Apply normalization rules
        validate: Validate extracted data
        log_callback: Optional callback for log messages

    Returns:
        Extraction result dictionary
    """
    extractor = MultiAgentTemplateExtractor(template_path)
    return extractor.extract(pdf_path, llm_model, normalize, validate, log_callback)


# Command line usage
if __name__ == "__main__":
    import sys
    from dotenv import load_dotenv
    load_dotenv()

    if len(sys.argv) < 3:
        print("Usage: python multiagent_template_extractor.py <pdf_path> <template_path> [llm_model]")
        print("\nExample:")
        print("  python multiagent_template_extractor.py module.pdf templates/pv_module_template.json")
        print("  python multiagent_template_extractor.py inverter.pdf templates/pv_inverter_template.json")
        sys.exit(1)

    pdf_path = sys.argv[1]
    template_path = sys.argv[2]
    llm_model = sys.argv[3] if len(sys.argv) > 3 else None

    result = extract_from_pdf(pdf_path, template_path, llm_model)

    print(f"\nSuccess: {result['success']}")
    if result['success']:
        print(f"Items extracted: {len(result['data'])}")

        for i, item in enumerate(result['data']):
            print(f"\n  Item {i+1}:")
            for key, value in list(item.items())[:8]:
                print(f"    {key}: {value}")
            if len(item) > 8:
                print(f"    ... and {len(item) - 8} more fields")
    else:
        print(f"Error: {result['error']}")