# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Université du Québec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

"""Parse downloaded filings into sections and chunks, in bounded batches.

    python scripts/parse_filings.py --max-workers 26 --batch-size 1500 --no-json

Only ``parse_status = 'pending'`` filings are processed, so the script can be
re-run after an interruption.  Results of each batch are written to DuckDB
before the next batch starts (memory stays bounded).
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.parsing import FilingParser
from src.storage.database import DatabaseManager
from src.utils.config import get_parsing_config
from src.utils.logging import setup_logging

logger = setup_logging("parse_filings")

SECTION_COLS = ["section_id", "accession_number", "section_name", "item_number",
                "section_order", "text", "word_count", "char_count"]
CHUNK_COLS = ["chunk_id", "section_id", "accession_number", "chunk_text", "chunk_order",
              "word_count", "char_count", "start_char", "end_char"]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Parse downloaded SEC filings into sections and chunks")
    parser.add_argument("--form-type", type=str, default=None, help="Filter by form type (e.g., 10-K)")
    parser.add_argument("--max-workers", type=int, default=4, help="Number of parallel workers (default: 4)")
    parser.add_argument("--batch-size", type=int, default=1000, help="Filings per batch (default: 1000)")
    parser.add_argument("--accession", type=str, default=None, help="Parse a specific filing by accession number")
    parser.add_argument("--no-json", action="store_true", help="Do not write per-filing JSON to data/processed")
    parser.add_argument("--limit", type=int, default=None, help="Stop after N filings (testing)")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    db = DatabaseManager()
    db.initialize()

    if args.accession:
        filing_df = db.fetch_df(
            "SELECT accession_number, form_type, local_path FROM filings WHERE accession_number = ?",
            [args.accession])
    else:
        conditions = ["download_status = 'complete'", "parse_status = 'pending'"]
        params: list = []
        if args.form_type:
            conditions.append("form_type = ?")
            params.append(args.form_type)
        where = " AND ".join(conditions)
        sql = (f"SELECT accession_number, form_type, local_path FROM filings WHERE {where} "
               f"ORDER BY filing_date")
        if args.limit:
            sql += f" LIMIT {int(args.limit)}"
        filing_df = db.fetch_df(sql, params if params else None)

    if filing_df.empty:
        logger.info("No filings to parse")
        db.close()
        return
    logger.info("Found %d filings to parse", len(filing_df))

    file_list, missing = [], []
    for row in filing_df.itertuples(index=False):
        lp = row.local_path or ""
        if not lp or not Path(lp).exists():
            missing.append(row.accession_number)
            continue
        file_list.append({"file_path": Path(lp), "form_type": row.form_type,
                          "accession_number": row.accession_number})
    if missing:
        logger.warning("%d filings have no local file; marking parse_status='missing'", len(missing))
        for acc in missing:
            db.update_filing_status(acc, "parse_status", "missing")

    config = dict(get_parsing_config())
    if args.no_json:
        config["save_json"] = False
    parser = FilingParser(config=config)

    total = len(file_list)
    succeeded = failed = 0
    t0 = time.time()
    for start in range(0, total, args.batch_size):
        batch = file_list[start:start + args.batch_size]
        results = parser.parse_batch(batch, max_workers=args.max_workers)

        sections, chunks, ok_acc, err_acc = [], [], [], []
        for r in results:
            if r.get("status") != "success":
                err_acc.append(r["accession_number"])
                continue
            sections.extend({k: s.get(k) for k in SECTION_COLS} for s in r["sections"])
            chunks.extend({k: c.get(k) for k in CHUNK_COLS} for c in r["chunks"])
            ok_acc.append(r["accession_number"])
        del results

        try:
            db.conn.execute("BEGIN")
            if sections:
                db.insert_records_df("filing_sections", sections)
            if chunks:
                db.insert_records_df("filing_chunks", chunks)
            if ok_acc:
                db.conn.register("_ok", pd.DataFrame({"accession_number": ok_acc}))
                db.conn.execute("UPDATE filings SET parse_status='complete', updated_at=current_timestamp "
                                "WHERE accession_number IN (SELECT accession_number FROM _ok)")
                db.conn.unregister("_ok")
            if err_acc:
                db.conn.register("_err", pd.DataFrame({"accession_number": err_acc}))
                db.conn.execute("UPDATE filings SET parse_status='error', updated_at=current_timestamp "
                                "WHERE accession_number IN (SELECT accession_number FROM _err)")
                db.conn.unregister("_err")
            db.conn.execute("COMMIT")
        except Exception:
            db.conn.execute("ROLLBACK")
            logger.error("Failed to store batch starting at %d", start, exc_info=True)
            raise
        succeeded += len(ok_acc)
        failed += len(err_acc)
        el = time.time() - t0
        done = min(start + args.batch_size, total)
        logger.info("batch done %d/%d filings | +%d sections +%d chunks | ok=%d err=%d | %.0fs (%.1f filings/s)",
                    done, total, len(sections), len(chunks), succeeded, failed, el, done / max(el, 1))

    logger.info("Parsing complete: %d succeeded, %d failed", succeeded, failed)
    db.close()


if __name__ == "__main__":
    main()