import json
import sqlite3
from datetime import datetime
from pathlib import Path

RESULTS_DIR = Path(__file__).parent.parent / "results"


def get_db(engagement: str) -> sqlite3.Connection:
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    db_path = RESULTS_DIR / f"{engagement}.db"
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    _init(conn)
    return conn


def _init(conn: sqlite3.Connection) -> None:
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS runs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            started_at TEXT NOT NULL,
            finished_at TEXT,
            countries TEXT,
            targets TEXT,
            total_hosts INTEGER DEFAULT 0,
            total_matches INTEGER DEFAULT 0
        );

        CREATE TABLE IF NOT EXISTS hits (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            run_id INTEGER NOT NULL,
            ip TEXT NOT NULL,
            port INTEGER NOT NULL,
            country_code TEXT,
            target_name TEXT NOT NULL,
            tags TEXT,
            probe_type TEXT,
            confidence TEXT,
            version TEXT,
            nmap_service TEXT,
            nmap_banner TEXT,
            found_at TEXT NOT NULL,
            FOREIGN KEY (run_id) REFERENCES runs(id)
        );

        CREATE INDEX IF NOT EXISTS idx_hits_ip ON hits(ip);
        CREATE INDEX IF NOT EXISTS idx_hits_target ON hits(target_name);
        CREATE INDEX IF NOT EXISTS idx_hits_country ON hits(country_code);
        CREATE INDEX IF NOT EXISTS idx_hits_run ON hits(run_id);
    """)
    conn.commit()


def start_run(conn: sqlite3.Connection, countries: list, targets: list) -> int:
    cur = conn.execute(
        "INSERT INTO runs (started_at, countries, targets) VALUES (?, ?, ?)",
        (datetime.utcnow().isoformat(), json.dumps(countries), json.dumps(targets)),
    )
    conn.commit()
    return cur.lastrowid


def finish_run(conn: sqlite3.Connection, run_id: int, total_hosts: int, total_matches: int) -> None:
    conn.execute(
        "UPDATE runs SET finished_at=?, total_hosts=?, total_matches=? WHERE id=?",
        (datetime.utcnow().isoformat(), total_hosts, total_matches, run_id),
    )
    conn.commit()


def insert_hit(conn: sqlite3.Connection, run_id: int, hit: dict) -> None:
    conn.execute(
        """INSERT INTO hits
           (run_id, ip, port, country_code, target_name, tags, probe_type,
            confidence, version, nmap_service, nmap_banner, found_at)
           VALUES (?,?,?,?,?,?,?,?,?,?,?,?)""",
        (
            run_id,
            hit["ip"],
            hit["port"],
            hit.get("country_code", ""),
            hit["target_name"],
            json.dumps(hit.get("tags", [])),
            hit.get("probe_type", ""),
            hit.get("confidence", ""),
            hit.get("version", ""),
            hit.get("nmap_service", ""),
            hit.get("nmap_banner", ""),
            datetime.utcnow().isoformat(),
        ),
    )
    conn.commit()


def export_json(conn: sqlite3.Connection, run_id: int | None = None) -> list[dict]:
    if run_id:
        rows = conn.execute("SELECT * FROM hits WHERE run_id=?", (run_id,)).fetchall()
    else:
        rows = conn.execute("SELECT * FROM hits").fetchall()
    return [dict(r) for r in rows]


def export_csv_lines(conn: sqlite3.Connection, run_id: int | None = None) -> list[str]:
    header = "ip,port,country_code,target_name,confidence,version,nmap_service,nmap_banner,found_at"
    if run_id:
        rows = conn.execute("SELECT * FROM hits WHERE run_id=?", (run_id,)).fetchall()
    else:
        rows = conn.execute("SELECT * FROM hits ORDER BY found_at").fetchall()
    lines = [header]
    for r in rows:
        lines.append(",".join([
            str(r["ip"]), str(r["port"]), r["country_code"] or "",
            r["target_name"], r["confidence"] or "", r["version"] or "",
            r["nmap_service"] or "", r["nmap_banner"] or "", r["found_at"],
        ]))
    return lines