"""TTP (Techniques, Tactics, Procedures) extractor for Mosaic."""
import logging
from typing import Optional

from .base import BaseExtractor
from storage.db import MosaicDB
from storage.models import TTP
from analysis.prompts import TTP_SYSTEM_PROMPT, TTP_EXTRACTION_PROMPT

logger = logging.getLogger(__name__)


class TTPExtractor(BaseExtractor):
    EXTRACTOR_NAME = "ttp"

    def get_system_prompt(self) -> str:
        return TTP_SYSTEM_PROMPT

    def get_extraction_prompt(self, text: str) -> str:
        return TTP_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:
                ttp = TTP(
                    technique=item.get('technique', ''),
                    category=item.get('category', ''),
                    description=item.get('description', ''),
                    actor=item.get('actor'),
                    mitre_id=item.get('mitre_id'),
                    source_doc_ids=[doc_id],
                    source_context=item.get('source_context', '')[:500],
                )
                ttp.compute_dedup_key()
                if self.db.insert_ttp(ttp):
                    count += 1
            except Exception as e:
                logger.warning("Failed to process TTP item: %s", e)
        return count