# ============================================================
# 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
# ============================================================

from __future__ import annotations

import json
import traceback
import uuid
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

from tqdm import tqdm

from src.parsing.chunker import TextChunker
from src.parsing.html_parser import HTMLCleaner
from src.parsing.section_extractor import SectionExtractor
from src.utils.config import get_parsing_config
from src.utils.constants import DATA_DIRS
from src.utils.logging import get_logger

logger = get_logger("parsing.pipeline")


def _parse_single(args: tuple[str, str, str, dict]) -> dict:
    """Top-level function so ProcessPoolExecutor can pickle it."""
    file_path_str, form_type, accession_number, config = args
    parser = FilingParser(config=config)
    return parser.parse_filing(Path(file_path_str), form_type, accession_number)


class FilingParser:
    def __init__(self, config: dict | None = None) -> None:
        self._config = config or get_parsing_config()
        chunking_cfg = self._config.get("chunking", {})
        self._cleaner = HTMLCleaner(config=self._config)
        self._extractor = SectionExtractor()
        self._chunker = TextChunker(
            chunk_size=chunking_cfg.get("chunk_size", 512),
            overlap=chunking_cfg.get("chunk_overlap", 64),
            min_chunk_size=chunking_cfg.get("min_chunk_size", 50),
        )
        parallel_cfg = self._config.get("parallel", {})
        self._max_workers: int = parallel_cfg.get("max_workers", 4)
        self._output_dir: Path = DATA_DIRS["processed"]
        self._save_json: bool = bool(self._config.get("save_json", True))

    def parse_filing(
        self, file_path: Path, form_type: str, accession_number: str
    ) -> dict:
        file_path = Path(file_path)
        logger.debug(
            "Parsing filing %s (%s) from %s", accession_number, form_type, file_path
        )

        try:
            clean_text = self._cleaner.clean_file(file_path)
        except Exception:
            logger.error(
                "Failed to clean %s: %s", accession_number, traceback.format_exc()
            )
            return self._error_result(accession_number, form_type, "clean_failed")

        if not clean_text:
            logger.warning("Empty text after cleaning: %s", accession_number)
            return self._error_result(accession_number, form_type, "empty_after_clean")

        sections = self._extractor.extract_sections(clean_text, form_type)
        if not sections and clean_text:
            # no Item markers found (e.g. exhibit-only 10-Q/10-K): keep the whole text as one
            # real section so the chunks satisfy the filing_chunks -> filing_sections foreign key
            sections = [{
                "section_id": str(uuid.uuid4()),
                "section_name": "full_document",
                "section_order": 0,
                "item_number": "N/A",
                "text": clean_text.strip(),
                "word_count": len(clean_text.split()),
                "char_count": len(clean_text),
            }]
        for section in sections:
            section["accession_number"] = accession_number

        chunks: list[dict] = []
        for section in sections:
            section_chunks = self._chunker.chunk_section(section)
            for chunk in section_chunks:
                chunk["accession_number"] = accession_number
            chunks.extend(section_chunks)

        result = {
            "accession_number": accession_number,
            "form_type": form_type,
            "clean_text": clean_text,
            "sections": sections,
            "chunks": chunks,
            "metadata": {
                "word_count": len(clean_text.split()),
                "char_count": len(clean_text),
                "section_count": len(sections),
                "chunk_count": len(chunks),
                "parse_timestamp": datetime.now(timezone.utc).isoformat(),
            },
            "status": "success",
        }

        if self._save_json:
            self._save_result(result)
        logger.debug(
            "Parsed %s: %d sections, %d chunks",
            accession_number,
            len(sections),
            len(chunks),
        )
        return result

    def parse_batch(
        self, file_list: list[dict], max_workers: int | None = None
    ) -> list[dict]:
        max_workers = max_workers or self._max_workers
        results: list[dict] = []
        failed = 0
        succeeded = 0

        tasks = [
            (
                str(item["file_path"]),
                item["form_type"],
                item["accession_number"],
                self._config,
            )
            for item in file_list
        ]

        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            futures = {
                executor.submit(_parse_single, task): task[2] for task in tasks
            }

            with tqdm(total=len(futures), desc="Parsing filings", unit="filing") as pbar:
                for future in as_completed(futures):
                    accession = futures[future]
                    try:
                        result = future.result()
                        results.append(result)
                        if result.get("status") == "success":
                            succeeded += 1
                        else:
                            failed += 1
                    except Exception:
                        logger.error(
                            "Batch parse failed for %s: %s",
                            accession,
                            traceback.format_exc(),
                        )
                        results.append(
                            self._error_result(accession, "unknown", "exception")
                        )
                        failed += 1
                    pbar.update(1)

        logger.info(
            "Batch complete: %d succeeded, %d failed out of %d total",
            succeeded,
            failed,
            len(file_list),
        )
        return results

    def _save_result(self, result: dict) -> None:
        self._output_dir.mkdir(parents=True, exist_ok=True)
        output_path = self._output_dir / f"{result['accession_number']}.json"

        serializable = {
            "accession_number": result["accession_number"],
            "form_type": result["form_type"],
            "sections": result["sections"],
            "chunks": result["chunks"],
            "metadata": result["metadata"],
            "status": result.get("status", "success"),
        }

        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(serializable, f, indent=2, ensure_ascii=False)

    def _error_result(
        self, accession_number: str, form_type: str, error_type: str
    ) -> dict:
        return {
            "accession_number": accession_number,
            "form_type": form_type,
            "clean_text": "",
            "sections": [],
            "chunks": [],
            "metadata": {
                "word_count": 0,
                "char_count": 0,
                "section_count": 0,
                "chunk_count": 0,
                "parse_timestamp": datetime.now(timezone.utc).isoformat(),
            },
            "status": f"error:{error_type}",
        }