"""Infrastructure indicator extractor for Mosaic (IPs, domains, C2, hashes)."""
import logging

from .base import BaseExtractor
from storage.models import Infrastructure
from analysis.prompts import INFRASTRUCTURE_SYSTEM_PROMPT, INFRASTRUCTURE_EXTRACTION_PROMPT

logger = logging.getLogger(__name__)


class InfrastructureExtractor(BaseExtractor):
    EXTRACTOR_NAME = "infrastructure"

    def __init__(self, *args, collection: str = "", doc_title: str = "",
                 doc_date: str = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.collection = collection
        self.doc_title = doc_title
        self.doc_date = doc_date

    def get_system_prompt(self) -> str:
        return INFRASTRUCTURE_SYSTEM_PROMPT

    def get_extraction_prompt(self, text: str) -> str:
        return INFRASTRUCTURE_EXTRACTION_PROMPT.format(document=text)

    def get_json_schema(self) -> dict:
        return {"properties": {"items": {"type": "array"}}, "required": ["items"]}

    def process_response(self, data: dict, doc_id: str) -> int:
        items = data.get('items', [])
        if isinstance(data, list):
            items = data
        count = 0
        for item in items:
            try:
                infra = Infrastructure(
                    indicator=item.get('indicator', ''),
                    indicator_type=item.get('indicator_type', ''),
                    context=item.get('context', ''),
                    source_doc_ids=[doc_id],
                    source_context=item.get('source_context', '')[:500],
                    collection=self.collection,
                    doc_title=self.doc_title,
                    doc_date=self.doc_date,
                )
                self.db.insert_infrastructure(infra)
                count += 1
            except Exception as e:
                logger.warning("Failed to process infrastructure item: %s", e)
        return count