# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
OSINT Investigation Persistence

SQLite-backed storage for OSINT investigations and findings.  Tracks the
full lifecycle from open → active → review → complete, with an entity
findings table that surfaces the same indicator (email, IP, domain, etc.)
appearing across multiple investigations — a high-value cross-case signal.

Database: ~/.familiar/data/osint/investigations.db

Schema:
    investigations — case record with target, justification, status
    findings       — discrete facts discovered (IP, email, subdomain, ...)

Tools exported:
    create_investigation, update_investigation_status, add_finding,
    list_investigations, get_investigation_findings, summarize_investigation
    (also: list_investigations, get_recent_findings for workspace loaders)
"""

import json
import logging
import sqlite3
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

_DB_DIR = Path.home() / ".familiar" / "data" / "osint"
_DB_PATH = _DB_DIR / "investigations.db"

# Valid investigation statuses (ordered for pipeline display)
VALID_STATUSES = ("open", "active", "review", "complete", "archived")

# Valid finding types
VALID_FINDING_TYPES = (
    "ip",
    "domain",
    "email",
    "subdomain",
    "credential",
    "person",
    "org",
    "phone",
    "certificate",
    "social_profile",
    "vulnerability",
    "netblock",
    "other",
)

# Valid risk levels
VALID_RISK_LEVELS = ("critical", "high", "medium", "low", "info")


# ---------------------------------------------------------------------------
# DB setup
# ---------------------------------------------------------------------------


def _get_db_path() -> Path:
    """Overridable for tests."""
    return _DB_PATH


def _get_conn() -> sqlite3.Connection:
    db_path = _get_db_path()
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    return conn


def _ensure_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS investigations (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            name        TEXT NOT NULL,
            target      TEXT NOT NULL,
            description TEXT,
            status      TEXT NOT NULL DEFAULT 'open',
            created_at  REAL NOT NULL,
            updated_at  REAL NOT NULL
        );

        CREATE TABLE IF NOT EXISTS findings (
            id               INTEGER PRIMARY KEY AUTOINCREMENT,
            investigation_id INTEGER NOT NULL REFERENCES investigations(id),
            finding_type     TEXT NOT NULL,
            value            TEXT NOT NULL,
            source_tool      TEXT NOT NULL,
            risk_level       TEXT,
            notes            TEXT,
            raw_data         TEXT,
            created_at       REAL NOT NULL
        );

        CREATE INDEX IF NOT EXISTS idx_finding_inv
            ON findings(investigation_id, finding_type);

        CREATE INDEX IF NOT EXISTS idx_finding_type
            ON findings(finding_type, value);
        """
    )
    conn.commit()


# ---------------------------------------------------------------------------
# Investigation CRUD
# ---------------------------------------------------------------------------


def create_investigation(
    name: str,
    target: str,
    description: str = "",
) -> Dict[str, Any]:
    """
    Create a new OSINT investigation case.

    Args:
        name: Short human-readable label for the investigation.
        target: Primary indicator being investigated (domain, IP, email, person, etc.).
        description: Optional longer context or justification for the investigation.

    Returns:
        Dict with investigation id, name, target, status, created_at.
    """
    if not name or not target:
        return {"error": "name and target are required"}

    now = time.time()
    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            cur = conn.execute(
                """
                INSERT INTO investigations (name, target, description, status, created_at, updated_at)
                VALUES (?, ?, ?, 'open', ?, ?)
                """,
                (name, target, description or "", now, now),
            )
            inv_id = cur.lastrowid
            conn.commit()
        return {
            "id": inv_id,
            "name": name,
            "target": target,
            "description": description,
            "status": "open",
            "created_at": now,
            "updated_at": now,
        }
    except Exception as e:
        logger.error("create_investigation failed: %s", e)
        return {"error": str(e)}


def update_investigation_status(
    investigation_id: int,
    status: str,
) -> Dict[str, Any]:
    """
    Update the pipeline status of an investigation.

    Args:
        investigation_id: ID of the investigation to update.
        status: New status — one of: open, active, review, complete, archived.

    Returns:
        Dict with id, status, updated_at, or error.
    """
    if status not in VALID_STATUSES:
        return {"error": f"Invalid status '{status}'. Must be one of: {', '.join(VALID_STATUSES)}"}

    now = time.time()
    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            cur = conn.execute(
                "UPDATE investigations SET status=?, updated_at=? WHERE id=?",
                (status, now, investigation_id),
            )
            conn.commit()
            if cur.rowcount == 0:
                return {"error": f"Investigation {investigation_id} not found"}
        return {"id": investigation_id, "status": status, "updated_at": now}
    except Exception as e:
        logger.error("update_investigation_status failed: %s", e)
        return {"error": str(e)}


def add_finding(
    investigation_id: int,
    finding_type: str,
    value: str,
    source_tool: str,
    risk_level: str = "info",
    notes: str = "",
    raw_data: Optional[Dict] = None,
) -> Dict[str, Any]:
    """
    Add a discrete finding to an investigation.

    Args:
        investigation_id: ID of the parent investigation.
        finding_type: Category of the finding — ip, domain, email, subdomain,
            credential, person, org, phone, certificate, social_profile,
            vulnerability, netblock, or other.
        value: The actual indicator value (e.g. "192.0.2.1", "user@example.com").
        source_tool: Which tool discovered this finding (e.g. "spiderfoot",
            "theHarvester", "shodan", "manual").
        risk_level: Severity — critical, high, medium, low, info. Default: info.
        notes: Optional analyst notes about this finding.
        raw_data: Optional raw JSON payload from the source tool.

    Returns:
        Dict with finding id, investigation_id, finding_type, value, risk_level.
    """
    if finding_type not in VALID_FINDING_TYPES:
        finding_type = "other"
    if risk_level not in VALID_RISK_LEVELS:
        risk_level = "info"

    now = time.time()
    raw_json = json.dumps(raw_data) if raw_data else None

    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            cur = conn.execute(
                """
                INSERT INTO findings
                    (investigation_id, finding_type, value, source_tool,
                     risk_level, notes, raw_data, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    investigation_id,
                    finding_type,
                    value,
                    source_tool,
                    risk_level,
                    notes or "",
                    raw_json,
                    now,
                ),
            )
            finding_id = cur.lastrowid
            # Touch parent updated_at
            conn.execute(
                "UPDATE investigations SET updated_at=? WHERE id=?",
                (now, investigation_id),
            )
            conn.commit()
        return {
            "id": finding_id,
            "investigation_id": investigation_id,
            "finding_type": finding_type,
            "value": value,
            "source_tool": source_tool,
            "risk_level": risk_level,
            "created_at": now,
        }
    except Exception as e:
        logger.error("add_finding failed: %s", e)
        return {"error": str(e)}


# ---------------------------------------------------------------------------
# Investigation queries
# ---------------------------------------------------------------------------


def list_investigations(
    status_filter: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """
    List investigations, optionally filtered by status.

    Args:
        status_filter: One of open, active, review, complete, archived.
            Pass None to list all.

    Returns:
        List of investigation dicts with finding_count included.
    """
    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            if status_filter:
                rows = conn.execute(
                    """
                    SELECT i.*, COUNT(f.id) AS finding_count
                    FROM investigations i
                    LEFT JOIN findings f ON f.investigation_id = i.id
                    WHERE i.status = ?
                    GROUP BY i.id
                    ORDER BY i.updated_at DESC
                    """,
                    (status_filter,),
                ).fetchall()
            else:
                rows = conn.execute(
                    """
                    SELECT i.*, COUNT(f.id) AS finding_count
                    FROM investigations i
                    LEFT JOIN findings f ON f.investigation_id = i.id
                    GROUP BY i.id
                    ORDER BY i.updated_at DESC
                    """
                ).fetchall()
        return [dict(row) for row in rows]
    except Exception as e:
        logger.error("list_investigations failed: %s", e)
        return []


def get_investigation_findings(
    investigation_id: int,
    type_filter: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Get all findings for an investigation, grouped by type.

    Args:
        investigation_id: ID of the investigation.
        type_filter: Optional finding_type to filter (e.g. "email", "ip").

    Returns:
        Dict with investigation_id, findings (list), type_groups (dict),
        and risk_groups (dict).
    """
    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            if type_filter:
                rows = conn.execute(
                    """
                    SELECT * FROM findings
                    WHERE investigation_id=? AND finding_type=?
                    ORDER BY risk_level, created_at DESC
                    """,
                    (investigation_id, type_filter),
                ).fetchall()
            else:
                rows = conn.execute(
                    """
                    SELECT * FROM findings
                    WHERE investigation_id=?
                    ORDER BY risk_level, created_at DESC
                    """,
                    (investigation_id,),
                ).fetchall()

        findings = [dict(row) for row in rows]

        type_groups: Dict[str, List] = {}
        risk_groups: Dict[str, int] = {}
        for f in findings:
            ft = f.get("finding_type", "other")
            type_groups.setdefault(ft, []).append(f)
            rl = f.get("risk_level", "info")
            risk_groups[rl] = risk_groups.get(rl, 0) + 1

        return {
            "investigation_id": investigation_id,
            "findings": findings,
            "total_count": len(findings),
            "type_groups": {k: len(v) for k, v in type_groups.items()},
            "risk_groups": risk_groups,
        }
    except Exception as e:
        logger.error("get_investigation_findings failed: %s", e)
        return {"investigation_id": investigation_id, "findings": [], "total_count": 0}


def summarize_investigation(investigation_id: int) -> Dict[str, Any]:
    """
    Generate a summary of an investigation with findings grouped by type and risk.

    Args:
        investigation_id: ID of the investigation to summarize.

    Returns:
        Dict with investigation metadata, finding counts by type and risk,
        and a list of notable (high/critical) findings.
    """
    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            inv_row = conn.execute(
                "SELECT * FROM investigations WHERE id=?", (investigation_id,)
            ).fetchone()
            if not inv_row:
                return {"error": f"Investigation {investigation_id} not found"}

            findings_rows = conn.execute(
                "SELECT * FROM findings WHERE investigation_id=? ORDER BY risk_level, created_at DESC",
                (investigation_id,),
            ).fetchall()

        investigation = dict(inv_row)
        findings = [dict(row) for row in findings_rows]

        by_type: Dict[str, int] = {}
        by_risk: Dict[str, int] = {}
        notable: List[Dict] = []

        for f in findings:
            ft = f.get("finding_type", "other")
            rl = f.get("risk_level", "info")
            by_type[ft] = by_type.get(ft, 0) + 1
            by_risk[rl] = by_risk.get(rl, 0) + 1
            if rl in ("critical", "high"):
                notable.append(f)

        return {
            "investigation": investigation,
            "finding_count": len(findings),
            "by_type": by_type,
            "by_risk": by_risk,
            "notable_findings": notable[:20],  # cap for readability
        }
    except Exception as e:
        logger.error("summarize_investigation failed: %s", e)
        return {"error": str(e)}


# ---------------------------------------------------------------------------
# Cross-investigation entity deduplication (workspace loader helper)
# ---------------------------------------------------------------------------


def get_recent_findings(limit: int = 20) -> List[Dict[str, Any]]:
    """
    Return the most recent findings across all investigations.

    Also flags (finding_type, value) pairs that appear in more than one
    investigation — a strong cross-case signal.

    Args:
        limit: Maximum rows to return. Default 20.

    Returns:
        List of finding dicts with added ``cross_case_count`` field.
    """
    try:
        with _get_conn() as conn:
            _ensure_schema(conn)
            # Count cross-investigation occurrences per (type, value) pair
            cross_counts = {
                (row[0], row[1]): row[2]
                for row in conn.execute(
                    """
                    SELECT finding_type, value, COUNT(DISTINCT investigation_id) AS n
                    FROM findings
                    GROUP BY finding_type, value
                    HAVING n > 1
                    """
                ).fetchall()
            }

            rows = conn.execute(
                """
                SELECT f.*, i.name AS investigation_name, i.target AS investigation_target
                FROM findings f
                JOIN investigations i ON i.id = f.investigation_id
                ORDER BY f.created_at DESC
                LIMIT ?
                """,
                (limit,),
            ).fetchall()

        result = []
        for row in rows:
            d = dict(row)
            d["cross_case_count"] = cross_counts.get((d.get("finding_type"), d.get("value")), 1)
            result.append(d)
        return result
    except Exception as e:
        logger.error("get_recent_findings failed: %s", e)
        return []
