#!/usr/bin/env python3
"""SQLite-first UAT/VAPT result database helper.

No third-party Python packages are required.  It can:
  - create schema
  - import generated Playwright matrix as test_cases
  - create a test run
  - record every Playwright test result immediately
  - finish a run and update counts
  - export latest run to JSON and a simple XLSX workbook
"""
import argparse
import csv
import datetime as dt
import html
import json
import os
import re
import sqlite3
import sys
import zipfile
from pathlib import Path

DB_DEFAULT = os.environ.get("SQLITE_RESULTS_DB", "results/uat_vapt_results.sqlite")
MATRIX_DEFAULT = os.environ.get("TEST_MATRIX_JSON", "data/generated-test-matrix.json")
EXTRA_CASES_DEFAULT = os.environ.get("EXTRA_TEST_CASES_JSON", "data/v2_5_test_cases.json,data/v2_5_1_test_cases.json")


def now_iso():
    return dt.datetime.now(dt.timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def connect(db_path: str):
    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=30000")
    return conn


def init_db(conn):
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS test_cases (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            source_id TEXT,
            variant_id TEXT UNIQUE,
            sheet TEXT,
            module TEXT,
            role TEXT,
            scenario TEXT,
            steps TEXT,
            expected TEXT,
            priority TEXT,
            type TEXT,
            route TEXT,
            variant_name TEXT,
            variant_type TEXT,
            automation TEXT,
            manual_reason TEXT,
            is_active INTEGER DEFAULT 1,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP,
            updated_at TEXT
        );

        CREATE TABLE IF NOT EXISTS test_runs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            run_name TEXT,
            base_url TEXT,
            status TEXT,
            started_at TEXT DEFAULT CURRENT_TIMESTAMP,
            finished_at TEXT,
            total_cases INTEGER DEFAULT 0,
            executed_count INTEGER DEFAULT 0,
            passed_count INTEGER DEFAULT 0,
            failed_count INTEGER DEFAULT 0,
            skipped_count INTEGER DEFAULT 0,
            timedout_count INTEGER DEFAULT 0,
            interrupted_count INTEGER DEFAULT 0,
            not_run_count INTEGER DEFAULT 0
        );

        CREATE TABLE IF NOT EXISTS test_results (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            run_id INTEGER NOT NULL,
            variant_id TEXT,
            test_title TEXT,
            title_path TEXT,
            project_name TEXT,
            file TEXT,
            line INTEGER,
            column INTEGER,
            status TEXT,
            expected_status TEXT,
            duration_ms INTEGER,
            retry INTEGER DEFAULT 0,
            error_message TEXT,
            error_stack TEXT,
            stdout TEXT,
            stderr TEXT,
            started_at TEXT,
            finished_at TEXT,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY(run_id) REFERENCES test_runs(id),
            FOREIGN KEY(variant_id) REFERENCES test_cases(variant_id)
        );

        CREATE TABLE IF NOT EXISTS test_attachments (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            result_id INTEGER NOT NULL,
            run_id INTEGER NOT NULL,
            variant_id TEXT,
            attachment_name TEXT,
            content_type TEXT,
            file_path TEXT,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY(result_id) REFERENCES test_results(id),
            FOREIGN KEY(run_id) REFERENCES test_runs(id)
        );

        CREATE INDEX IF NOT EXISTS idx_test_cases_variant_id ON test_cases(variant_id);
        CREATE INDEX IF NOT EXISTS idx_test_results_run_id ON test_results(run_id);
        CREATE INDEX IF NOT EXISTS idx_test_results_variant_id ON test_results(variant_id);
        CREATE INDEX IF NOT EXISTS idx_test_results_status ON test_results(status);
        CREATE INDEX IF NOT EXISTS idx_test_attachments_run_id ON test_attachments(run_id);
        """
    )
    conn.commit()


def norm(v):
    return "" if v is None else str(v)


def _load_case_rows(path: str):
    if not path or not Path(path).exists():
        return []
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def import_cases(conn, matrix_path: str):
    """Import primary generated matrix plus optional v2.5/additional test-case JSON.

    matrix_path may be a comma-separated list. In addition, EXTRA_TEST_CASES_JSON
    defaults to data/v2_5_test_cases.json when present.
    """
    init_db(conn)
    paths = []
    for part in str(matrix_path or "").split(','):
        part = part.strip()
        if part:
            paths.append(part)
    extra = EXTRA_CASES_DEFAULT
    for part in str(extra or "").split(','):
        part = part.strip()
        if part and part not in paths:
            paths.append(part)

    sql = """
        INSERT INTO test_cases (
            source_id, variant_id, sheet, module, role, scenario, steps, expected,
            priority, type, route, variant_name, variant_type, automation, manual_reason,
            is_active, updated_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?)
        ON CONFLICT(variant_id) DO UPDATE SET
            source_id=excluded.source_id,
            sheet=excluded.sheet,
            module=excluded.module,
            role=excluded.role,
            scenario=excluded.scenario,
            steps=excluded.steps,
            expected=excluded.expected,
            priority=excluded.priority,
            type=excluded.type,
            route=excluded.route,
            variant_name=excluded.variant_name,
            variant_type=excluded.variant_type,
            automation=excluded.automation,
            manual_reason=excluded.manual_reason,
            is_active=1,
            updated_at=excluded.updated_at
    """
    ts = now_iso()
    count = 0
    for path in paths:
        rows = _load_case_rows(path)
        for r in rows:
            variant_id = r.get("variantId") or r.get("variant_id") or ""
            if not variant_id:
                continue
            conn.execute(sql, (
                norm(r.get("sourceId") or r.get("source_id")),
                norm(variant_id),
                norm(r.get("sheet")),
                norm(r.get("module")),
                norm(r.get("role")),
                norm(r.get("scenario")),
                norm(r.get("steps")),
                norm(r.get("expected")),
                norm(r.get("priority")),
                norm(r.get("type")),
                norm(r.get("route")),
                norm(r.get("variantName") or r.get("variant_name")),
                norm(r.get("variantType") or r.get("variant_type")),
                norm(r.get("automation")),
                norm(r.get("manualReason") or r.get("manual_reason")),
                ts,
            ))
            count += 1
    conn.commit()
    return count

def begin_run(conn, run_name: str, base_url: str):
    init_db(conn)
    total_cases = conn.execute("SELECT COUNT(*) AS c FROM test_cases WHERE is_active=1").fetchone()["c"]
    cur = conn.execute(
        "INSERT INTO test_runs (run_name, base_url, status, started_at, total_cases) VALUES (?, ?, ?, ?, ?)",
        (run_name, base_url, "running", now_iso(), total_cases),
    )
    conn.commit()
    return cur.lastrowid


def record_result(conn, record: dict):
    init_db(conn)
    run_id = int(record.get("run_id") or record.get("runId") or 0)
    if not run_id:
        raise ValueError("run_id is required")
    variant_id = record.get("variant_id") or record.get("variantId") or extract_variant_id(" ".join(record.get("title_path") or record.get("titlePath") or []))
    if isinstance(record.get("title_path"), list):
        title_path = json.dumps(record.get("title_path"), ensure_ascii=False)
    elif isinstance(record.get("titlePath"), list):
        title_path = json.dumps(record.get("titlePath"), ensure_ascii=False)
    else:
        title_path = norm(record.get("title_path") or record.get("titlePath"))

    err = record.get("error") or {}
    if isinstance(err, str):
        error_message, error_stack = err, ""
    else:
        error_message = norm(err.get("message"))
        error_stack = norm(err.get("stack"))

    loc = record.get("location") or {}
    cur = conn.execute(
        """
        INSERT INTO test_results (
            run_id, variant_id, test_title, title_path, project_name, file, line, column,
            status, expected_status, duration_ms, retry, error_message, error_stack,
            stdout, stderr, started_at, finished_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            run_id,
            norm(variant_id),
            norm(record.get("title") or record.get("test_title")),
            title_path,
            norm(record.get("project_name") or record.get("projectName")),
            norm(loc.get("file") or record.get("file")),
            int(loc.get("line") or record.get("line") or 0),
            int(loc.get("column") or record.get("column") or 0),
            norm(record.get("status")),
            norm(record.get("expected_status") or record.get("expectedStatus")),
            int(record.get("duration_ms") or record.get("duration") or 0),
            int(record.get("retry") or 0),
            error_message,
            error_stack,
            norm(record.get("stdout")),
            norm(record.get("stderr")),
            norm(record.get("started_at") or record.get("startedAt")),
            norm(record.get("finished_at") or record.get("finishedAt") or now_iso()),
        ),
    )
    result_id = cur.lastrowid
    for a in record.get("attachments") or []:
        conn.execute(
            """
            INSERT INTO test_attachments (result_id, run_id, variant_id, attachment_name, content_type, file_path)
            VALUES (?, ?, ?, ?, ?, ?)
            """,
            (result_id, run_id, norm(variant_id), norm(a.get("name")), norm(a.get("contentType") or a.get("content_type")), norm(a.get("path") or a.get("file_path"))),
        )
    conn.commit()
    return result_id


def finish_run(conn, run_id: int, status: str):
    init_db(conn)
    rows = conn.execute("SELECT status, COUNT(*) AS c FROM test_results WHERE run_id=? GROUP BY status", (run_id,)).fetchall()
    counts = {r["status"]: r["c"] for r in rows}
    executed = sum(counts.values())
    total_cases = conn.execute("SELECT COUNT(*) AS c FROM test_cases WHERE is_active=1").fetchone()["c"]
    not_run = max(total_cases - executed, 0)
    conn.execute(
        """
        UPDATE test_runs SET
            status=?, finished_at=?, total_cases=?, executed_count=?, passed_count=?, failed_count=?,
            skipped_count=?, timedout_count=?, interrupted_count=?, not_run_count=?
        WHERE id=?
        """,
        (
            status,
            now_iso(),
            total_cases,
            executed,
            counts.get("passed", 0),
            counts.get("failed", 0),
            counts.get("skipped", 0),
            counts.get("timedOut", 0) + counts.get("timedout", 0),
            counts.get("interrupted", 0),
            not_run,
            run_id,
        ),
    )
    conn.commit()
    return counts


def latest_run_id(conn):
    row = conn.execute("SELECT id FROM test_runs ORDER BY id DESC LIMIT 1").fetchone()
    return row["id"] if row else None


def merged_rows(conn, run_id: int):
    # Use the latest result per variant_id within the run. Retries become later rows; last row wins.
    sql = """
    WITH latest AS (
        SELECT tr.*
        FROM test_results tr
        JOIN (
            SELECT variant_id, MAX(id) AS max_id
            FROM test_results
            WHERE run_id=?
            GROUP BY variant_id
        ) x ON x.max_id = tr.id
    )
    SELECT
        tc.source_id AS sourceId,
        tc.sheet,
        tc.module,
        tc.role,
        tc.scenario,
        tc.steps,
        tc.expected,
        tc.priority,
        tc.type,
        tc.route,
        tc.variant_id AS variantId,
        tc.variant_name AS variantName,
        tc.variant_type AS variantType,
        tc.automation,
        COALESCE(latest.status, CASE WHEN tc.automation='manual' THEN 'manual' ELSE 'not_run' END) AS result,
        latest.duration_ms AS durationMs,
        latest.retry,
        latest.error_message AS errorMessage,
        latest.started_at AS startedAt,
        latest.finished_at AS finishedAt
    FROM test_cases tc
    LEFT JOIN latest ON latest.variant_id = tc.variant_id
    WHERE tc.is_active=1
    ORDER BY tc.sheet, tc.source_id, tc.variant_id
    """
    return [dict(r) for r in conn.execute(sql, (run_id,)).fetchall()]


def export_merged_json(conn, run_id: int, output_path: str):
    rows = merged_rows(conn, run_id)
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(rows, f, indent=2, ensure_ascii=False)
    return len(rows)


def export_csv(conn, run_id: int, output_path: str):
    rows = merged_rows(conn, run_id)
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", newline="", encoding="utf-8") as f:
        if rows:
            w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
            w.writeheader()
            w.writerows(rows)
    return len(rows)


def xlsx_col(n):
    s = ""
    while n:
        n, r = divmod(n - 1, 26)
        s = chr(65 + r) + s
    return s


def sheet_xml(rows):
    out = ['<?xml version="1.0" encoding="UTF-8" standalone="yes"?>', '<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main"><sheetData>']
    for r_idx, row in enumerate(rows, start=1):
        out.append(f'<row r="{r_idx}">')
        for c_idx, val in enumerate(row, start=1):
            ref = f"{xlsx_col(c_idx)}{r_idx}"
            if val is None:
                val = ""
            # Keep everything as inline string for maximum compatibility with no sharedStrings table.
            txt = html.escape(str(val), quote=False)
            out.append(f'<c r="{ref}" t="inlineStr"><is><t>{txt}</t></is></c>')
        out.append('</row>')
    out.append('</sheetData></worksheet>')
    return "".join(out)


def write_xlsx(path_out: str, sheets: dict):
    Path(path_out).parent.mkdir(parents=True, exist_ok=True)
    names = list(sheets.keys())
    with zipfile.ZipFile(path_out, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("[Content_Types].xml", """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>
""" + "".join(f'<Override PartName="/xl/worksheets/sheet{i}.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>' for i in range(1, len(names)+1)) + "</Types>")
        z.writestr("_rels/.rels", """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>
</Relationships>""")
        z.writestr("xl/workbook.xml", """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"><sheets>""" + "".join(f'<sheet name="{html.escape(name[:31])}" sheetId="{i}" r:id="rId{i}"/>' for i, name in enumerate(names, start=1)) + "</sheets></workbook>")
        z.writestr("xl/_rels/workbook.xml.rels", """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">""" + "".join(f'<Relationship Id="rId{i}" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" Target="worksheets/sheet{i}.xml"/>' for i in range(1, len(names)+1)) + "</Relationships>")
        for i, name in enumerate(names, start=1):
            z.writestr(f"xl/worksheets/sheet{i}.xml", sheet_xml(sheets[name]))


def rows_to_aoa(rows):
    if not rows:
        return []
    headers = list(rows[0].keys())
    return [headers] + [[r.get(h, "") for h in headers] for r in rows]


def export_xlsx(conn, run_id: int, output_path: str):
    run = conn.execute("SELECT * FROM test_runs WHERE id=?", (run_id,)).fetchone()
    summary = conn.execute("SELECT status, COUNT(*) AS count FROM test_results WHERE run_id=? GROUP BY status", (run_id,)).fetchall()
    results = merged_rows(conn, run_id)
    attachments = conn.execute(
        """
        SELECT tr.variant_id, tr.test_title, ta.attachment_name, ta.content_type, ta.file_path
        FROM test_attachments ta
        JOIN test_results tr ON tr.id = ta.result_id
        WHERE ta.run_id=?
        ORDER BY tr.variant_id, ta.attachment_name
        """,
        (run_id,),
    ).fetchall()
    failures = [r for r in results if r.get("result") not in ("passed", "manual", "skipped")]
    sheets = {
        "Run": rows_to_aoa([dict(run)] if run else []),
        "Summary": rows_to_aoa([dict(r) for r in summary]),
        "Failures": rows_to_aoa(failures),
        "Results": rows_to_aoa(results),
        "Attachments": rows_to_aoa([dict(r) for r in attachments]),
    }
    write_xlsx(output_path, sheets)
    return len(results)


def extract_variant_id(text: str):
    text = text or ""
    m = re.search(r"\[([^\]]+)\]", text)
    if m:
        return m.group(1)
    # Fallback for legacy/custom tests whose title starts with an ID.
    m = re.search(r"\b((?:V25-)?[A-Z][A-Z0-9]*(?:-[A-Z0-9]+)+)\b", text)
    return m.group(1) if m else ""


def main():
    p = argparse.ArgumentParser()
    p.add_argument("command", choices=["init", "import-cases", "begin-run", "record-result", "finish-run", "export-json", "export-csv", "export-xlsx", "summary"])
    p.add_argument("--db", default=DB_DEFAULT)
    p.add_argument("--matrix", default=MATRIX_DEFAULT)
    p.add_argument("--run-id", type=int)
    p.add_argument("--run-name", default=os.environ.get("TEST_RUN_NAME") or f"Run {now_iso()}")
    p.add_argument("--base-url", default=os.environ.get("BASE_URL", ""))
    p.add_argument("--status", default="completed")
    p.add_argument("--output", default="")
    args = p.parse_args()

    conn = connect(args.db)
    try:
        if args.command == "init":
            init_db(conn)
            print(json.dumps({"ok": True, "db": args.db}))
        elif args.command == "import-cases":
            count = import_cases(conn, args.matrix)
            print(json.dumps({"ok": True, "imported": count, "db": args.db}))
        elif args.command == "begin-run":
            run_id = begin_run(conn, args.run_name, args.base_url)
            print(json.dumps({"ok": True, "run_id": run_id}))
        elif args.command == "record-result":
            record = json.load(sys.stdin)
            result_id = record_result(conn, record)
            print(json.dumps({"ok": True, "result_id": result_id}))
        elif args.command == "finish-run":
            run_id = args.run_id or latest_run_id(conn)
            counts = finish_run(conn, run_id, args.status)
            print(json.dumps({"ok": True, "run_id": run_id, "counts": counts}))
        elif args.command == "export-json":
            run_id = args.run_id or latest_run_id(conn)
            output = args.output or "results/merged-results.json"
            count = export_merged_json(conn, run_id, output)
            print(json.dumps({"ok": True, "run_id": run_id, "rows": count, "output": output}))
        elif args.command == "export-csv":
            run_id = args.run_id or latest_run_id(conn)
            output = args.output or "results/UAT_VAPT_SQLite_Export.csv"
            count = export_csv(conn, run_id, output)
            print(json.dumps({"ok": True, "run_id": run_id, "rows": count, "output": output}))
        elif args.command == "export-xlsx":
            run_id = args.run_id or latest_run_id(conn)
            output = args.output or "results/UAT_VAPT_SQLite_Export.xlsx"
            count = export_xlsx(conn, run_id, output)
            print(json.dumps({"ok": True, "run_id": run_id, "rows": count, "output": output}))
        elif args.command == "summary":
            run_id = args.run_id or latest_run_id(conn)
            rows = conn.execute("SELECT status, COUNT(*) AS count FROM test_results WHERE run_id=? GROUP BY status", (run_id,)).fetchall()
            print(json.dumps({r["status"]: r["count"] for r in rows}, indent=2))
    finally:
        conn.close()


if __name__ == "__main__":
    main()