#!/usr/bin/env python3
"""
ransomware.live PRO API fetcher.
Polls for new ransomware victims and converts them to article dicts
compatible with the existing send_alert pipeline.
"""

import asyncio
import aiohttp
import sqlite3
import logging
from datetime import datetime, timezone, timedelta
from typing import Dict, List, Optional, Tuple
from urllib.parse import urlparse
from html import escape

from bot_config import SEEN_RETENTION_DAYS, ALLOWED_RANSOMWARE_COUNTRIES

logger = logging.getLogger(__name__)

API_BASE = "https://api-pro.ransomware.live"


class RansomwareFetcher:

    def __init__(self, api_key: str, seen_file: str = "seen_victims.db"):
        self.api_key = api_key
        self.seen_file = seen_file
        self._ensure_db()

    def _ensure_db(self):
        with sqlite3.connect(self.seen_file) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS seen_victims (
                    victim_id TEXT PRIMARY KEY,
                    seen_at   TEXT NOT NULL
                )
            """)
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_seen_victims_seen_at ON seen_victims (seen_at)"
            )
            conn.commit()

    def _is_seen(self, victim_id: str) -> bool:
        with sqlite3.connect(self.seen_file) as conn:
            return conn.execute(
                "SELECT 1 FROM seen_victims WHERE victim_id = ?", (victim_id,)
            ).fetchone() is not None

    def is_first_run(self) -> bool:
        """True if no victims have ever been recorded as seen (fresh DB)."""
        with sqlite3.connect(self.seen_file) as conn:
            count = conn.execute("SELECT COUNT(1) FROM seen_victims").fetchone()[0]
        return count == 0

    def mark_seen(self, victim_id: str):
        now = datetime.now(timezone.utc).isoformat()
        cutoff = (datetime.now(timezone.utc) - timedelta(days=SEEN_RETENTION_DAYS)).isoformat()
        with sqlite3.connect(self.seen_file) as conn:
            conn.execute("DELETE FROM seen_victims WHERE seen_at < ?", (cutoff,))
            conn.execute(
                "INSERT OR REPLACE INTO seen_victims (victim_id, seen_at) VALUES (?, ?)",
                (victim_id, now)
            )
            conn.commit()

    async def _fetch_month(
        self, session: aiohttp.ClientSession, year: int, month: int
    ) -> List[Dict]:
        url = f"{API_BASE}/victims/"
        params = {"year": year, "month": f"{month:02d}"}
        headers = {"X-Api-Key": self.api_key}
        try:
            async with session.get(url, params=params, headers=headers, timeout=aiohttp.ClientTimeout(total=30)) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    return data.get("victims", [])
                elif resp.status == 401:
                    logger.error("ransomware.live: invalid API key")
                elif resp.status == 429:
                    logger.warning("ransomware.live: rate limited, skipping this poll")
                else:
                    logger.error(f"ransomware.live: HTTP {resp.status}")
        except asyncio.TimeoutError:
            logger.error("ransomware.live: request timed out")
        except Exception as e:
            logger.error(f"ransomware.live: {e}")
        return []

    async def fetch_new_victims(
        self, session: aiohttp.ClientSession, initial_run: bool = False
    ) -> List[Dict]:
        now = datetime.now(timezone.utc)
        cutoff = now - timedelta(hours=48)

        months = [(now.year, now.month)]
        if now.day <= 2:
            prev = (now.replace(day=1) - timedelta(days=1))
            months.append((prev.year, prev.month))

        all_victims: List[Dict] = []
        for year, month in months:
            victims = await self._fetch_month(session, year, month)
            all_victims.extend(victims)

        # Only care about victims in these countries — the general global feed
        # is available directly on ransomware.live's own site.
        all_victims = [
            v for v in all_victims
            if (v.get("country") or "").strip().upper() in ALLOWED_RANSOMWARE_COUNTRIES
        ]

        new_victims = []
        for v in all_victims:
            discovered_str = v.get("discovered", "")
            if not discovered_str:
                continue
            try:
                discovered_dt = datetime.fromisoformat(discovered_str.replace("Z", "+00:00"))
            except ValueError:
                continue
            if discovered_dt.tzinfo is None:
                discovered_dt = discovered_dt.replace(tzinfo=timezone.utc)
            if discovered_dt < cutoff:
                continue

            vid = v.get("id", "")
            if not vid:
                continue
            if self._is_seen(vid):
                continue

            if initial_run:
                self.mark_seen(vid)
            else:
                new_victims.append(v)

        return new_victims

    @staticmethod
    def _extract_company(victim: Dict) -> str:
        raw = victim.get("victim", "")
        if raw and not raw.startswith("http"):
            return raw
        domain = victim.get("website", "")
        if domain:
            return domain
        if raw.startswith("http"):
            return urlparse(raw).netloc or raw
        return "Unknown"

    @staticmethod
    def to_article(victim: Dict) -> Dict:
        company  = RansomwareFetcher._extract_company(victim)
        group    = victim.get("group", "unknown")
        country  = victim.get("country", "")
        sector   = victim.get("activity", "")
        raw_desc = victim.get("description") or ""

        # Drop AI-generated placeholder descriptions
        description = "" if "[AI generated]" in raw_desc else raw_desc[:300]

        meta_parts = []
        if sector:  meta_parts.append(f"Sector: {sector}")
        if country: meta_parts.append(f"Country: {country}")
        full_desc = " | ".join(meta_parts)
        if description:
            full_desc = f"{full_desc}\n{description}" if full_desc else description

        discovered = victim.get("discovered", "")
        try:
            dt = datetime.fromisoformat(discovered.replace("Z", "+00:00"))
            published_human = dt.strftime("%Y-%m-%d %H:%M UTC")
        except (ValueError, AttributeError):
            published_human = discovered[:10] if discovered else "Unknown"

        screenshot = victim.get("screenshot") or None
        if screenshot and not screenshot.startswith("http"):
            screenshot = None

        return {
            "title":            f"{company} claimed by {group.title()}",
            "description":      full_desc,
            "url":              victim.get("permalink", ""),
            "published_human":  published_human,
            "source":           "ransomware.live",
            "category":         "ransomware",
            "feed_type":        "ransomware",
            "thumbnail":        screenshot,
            "cves":             [],
            "threat_actors":    [],
            "malware_families": [group.title()],
            "mitre_techniques": [],
        }