"""Named tool/exploit/implant extractor for Mosaic."""
import json
import logging
from pathlib import Path
from typing import Optional

from .base import BaseExtractor
from storage.db import MosaicDB
from storage.models import ExtractedTool
from analysis.prompts import TOOL_SYSTEM_PROMPT, TOOL_EXTRACTION_PROMPT

logger = logging.getLogger(__name__)

KNOWN_TOOLS_PATH = Path(__file__).parent.parent / 'data' / 'known_tools.json'


class ToolExtractor(BaseExtractor):
    EXTRACTOR_NAME = "tool"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.known_tools = self._load_known_tools()

    def _load_known_tools(self) -> list[str]:
        try:
            with open(KNOWN_TOOLS_PATH) as f:
                data = json.load(f)
            names = []
            for tool in data.get('tools', []):
                names.append(tool['name'].lower())
                for alias in tool.get('aliases', []):
                    names.append(alias.lower())
            return names
        except Exception as e:
            logger.warning("Could not load known tools: %s", e)
            return []

    def get_system_prompt(self) -> str:
        return TOOL_SYSTEM_PROMPT

    def get_extraction_prompt(self, text: str) -> str:
        # Include known tools as context
        known_hint = ', '.join(self.known_tools[:50]) if self.known_tools else 'none loaded'
        return TOOL_EXTRACTION_PROMPT.format(document=text, known_tools=known_hint)

    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:
                tool = ExtractedTool(
                    name=item.get('name', ''),
                    aliases=item.get('aliases', []),
                    description=item.get('description', ''),
                    capability=item.get('capability', ''),
                    target_platforms=item.get('target_platforms', []),
                    target_software=item.get('target_software', []),
                    cves=item.get('cves', []),
                    source_doc_ids=[doc_id],
                    source_context=item.get('source_context', '')[:500],
                    mitre_techniques=item.get('mitre_techniques', []),
                )
                tool.compute_dedup_key()
                if self.db.insert_extracted_tool(tool):
                    count += 1
            except Exception as e:
                logger.warning("Failed to process tool item: %s", e)
        return count