"""JSONL export/import for Mosaic extracted data."""
import json
import hashlib
import logging
from pathlib import Path
from datetime import datetime
from typing import Optional

from .db import MosaicDB

logger = logging.getLogger(__name__)

CATEGORIES = {
    'tools': ('extracted_tools', 'get_all_tools'),
    'ttps': ('ttps', 'get_all_ttps'),
    'tradecraft': ('tradecraft', 'get_all_tradecraft'),
    'entities': ('entities', 'get_all_entities'),
    'infrastructure': ('infrastructure', 'get_all_infrastructure'),
}


class JSONLStore:
    """Export and import extracted data as JSONL files with checksums."""

    def __init__(self, db: MosaicDB, output_dir: str = "output/extracted"):
        self.db = db
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def export(self, category: Optional[str] = None, source: Optional[str] = None) -> list[str]:
        """Export data to JSONL files. Returns list of created file paths."""
        categories = [category] if category else list(CATEGORIES.keys())
        created = []

        for cat in categories:
            if cat not in CATEGORIES:
                logger.warning("Unknown category: %s", cat)
                continue

            table_name, getter_name = CATEGORIES[cat]
            getter = getattr(self.db, getter_name)
            items = getter()

            if not items:
                logger.info("No data for category: %s", cat)
                continue

            timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
            filename = f"{cat}_{timestamp}.jsonl"
            filepath = self.output_dir / filename

            sha256 = hashlib.sha256()
            count = 0
            with open(filepath, 'w') as f:
                for item in items:
                    line = json.dumps(item.to_dict(), ensure_ascii=False) + '\n'
                    f.write(line)
                    sha256.update(line.encode())
                    count += 1

            # Write checksum file
            checksum_path = filepath.with_suffix('.jsonl.sha256')
            checksum_path.write_text(f"{sha256.hexdigest()}  {filename}\n")

            logger.info("Exported %d %s to %s", count, cat, filepath)
            created.append(str(filepath))

        return created

    def import_jsonl(self, filepath: str, category: str) -> int:
        """Import data from a JSONL file. Returns count of imported records."""
        from .models import ExtractedTool, TTP, Tradecraft, Entity, Infrastructure

        model_map = {
            'tools': ExtractedTool,
            'ttps': TTP,
            'tradecraft': Tradecraft,
            'entities': Entity,
            'infrastructure': Infrastructure,
        }
        insert_map = {
            'tools': self.db.insert_extracted_tool,
            'ttps': self.db.insert_ttp,
            'tradecraft': self.db.insert_tradecraft,
            'entities': self.db.insert_entity,
            'infrastructure': self.db.insert_infrastructure,
        }

        if category not in model_map:
            raise ValueError(f"Unknown category: {category}")

        model_cls = model_map[category]
        inserter = insert_map[category]
        count = 0

        with open(filepath, 'r') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                data = json.loads(line)
                # Deserialize JSON list fields
                obj = model_cls.from_row(data)
                inserter(obj)
                count += 1

        logger.info("Imported %d %s from %s", count, category, filepath)
        return count

    def verify_checksum(self, filepath: str) -> bool:
        """Verify a JSONL file against its checksum."""
        checksum_path = Path(filepath).with_suffix('.jsonl.sha256')
        if not checksum_path.exists():
            logger.warning("No checksum file found for %s", filepath)
            return False

        expected = checksum_path.read_text().split()[0]
        sha256 = hashlib.sha256()
        with open(filepath, 'rb') as f:
            for chunk in iter(lambda: f.read(8192), b''):
                sha256.update(chunk)

        actual = sha256.hexdigest()
        if actual != expected:
            logger.error("Checksum mismatch for %s: expected %s, got %s",
                        filepath, expected, actual)
            return False
        return True