from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path(
    __file__
).resolve().parents[1]

if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(
        0,
        str(PROJECT_ROOT),
    )


from evaluation.metrics import (
    aggregate_by_category,
    aggregate_by_difficulty,
    aggregate_metrics,
    evaluate_question_mode,
)
from evaluation.retrieval_runner import (
    EVALUATION_MODES,
    count_splits,
    filter_questions_by_split,
    load_index_document_paths,
    load_questions,
    retrieve_all_modes,
    validate_dataset,
)
from scripts.common import DB_FILE


def print_summary(
    summary: dict[
        str,
        dict[str, Any],
    ],
    *,
    split: str,
    selected_count: int,
    total_count: int,
) -> None:
    print()
    print(
        "Retrieval evaluation"
    )
    print(
        "=" * 78
    )

    print(
        f"Split: {split}"
    )

    print(
        f"Questions: "
        f"{selected_count}/{total_count}"
    )

    print(
        "-" * 78
    )

    header = (
        f"{'Mode':<10}"
        f"{'Questions':>10}"
        f"{'Hit@1':>10}"
        f"{'Hit@3':>10}"
        f"{'Hit@5':>10}"
        f"{'MRR':>10}"
        f"{'Recall@5':>12}"
    )

    print(
        header
    )

    print(
        "-" * 78
    )

    for mode in EVALUATION_MODES:
        metrics = summary[
            mode
        ]

        print(
            f"{mode:<10}"
            f"{metrics['questions']:>10}"
            f"{metrics['hit_at_1']:>10.3f}"
            f"{metrics['hit_at_3']:>10.3f}"
            f"{metrics['hit_at_5']:>10.3f}"
            f"{metrics['mrr']:>10.3f}"
            f"{metrics['recall_at_5']:>12.3f}"
        )

    print(
        "=" * 78
    )
    print()


def save_json_results(
    path: Path,
    payload: dict[str, Any],
) -> None:
    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    with path.open(
        "w",
        encoding="utf-8",
    ) as file:
        json.dump(
            payload,
            file,
            ensure_ascii=False,
            indent=2,
        )

        file.write(
            "\n"
        )


def save_csv_results(
    path: Path,
    rows: list[
        dict[str, Any]
    ],
) -> None:
    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    fieldnames = [
        "id",
        "split",
        "category",
        "difficulty",
        "mode",
        "question",
        "first_relevant_rank",
        "reciprocal_rank",
        "hit_at_1",
        "hit_at_3",
        "hit_at_5",
        "recall_at_1",
        "recall_at_3",
        "recall_at_5",
        "expected_documents",
        "top_documents",
        "top_source_urls",
    ]

    with path.open(
        "w",
        encoding="utf-8",
        newline="",
    ) as file:
        writer = csv.DictWriter(
            file,
            fieldnames=fieldnames,
        )

        writer.writeheader()

        for row in rows:
            csv_row = dict(
                row
            )

            for key in (
                "expected_documents",
                "top_documents",
                "top_source_urls",
            ):
                csv_row[
                    key
                ] = json.dumps(
                    csv_row.get(
                        key,
                        [],
                    ),
                    ensure_ascii=False,
                )

            writer.writerow(
                {
                    key: csv_row.get(
                        key
                    )
                    for key in fieldnames
                }
            )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Porovnanie FTS5, embeddingového "
            "a hybridného retrievalu."
        )
    )

    parser.add_argument(
        "--questions",
        type=Path,
        default=(
            PROJECT_ROOT
            / "evaluation"
            / "questions.json"
        ),
        help=(
            "Cesta k questions.json"
        ),
    )

    parser.add_argument(
        "--db",
        type=Path,
        default=DB_FILE,
        help=(
            "Cesta k SQLite indexu"
        ),
    )

    parser.add_argument(
        "--output-dir",
        type=Path,
        default=(
            PROJECT_ROOT
            / "evaluation"
            / "results"
        ),
        help=(
            "Adresár pre výsledky"
        ),
    )

    parser.add_argument(
        "--split",
        choices=(
            "dev",
            "test",
            "all",
        ),
        default="dev",
        help=(
            "Časť datasetu: "
            "dev = ladenie, "
            "test = finálne hodnotenie, "
            "all = celý dataset. "
            "Predvolené je dev."
        ),
    )

    parser.add_argument(
        "--limit",
        type=int,
        default=5,
        help=(
            "Počet dokumentov použitých "
            "pri evaluácii. Minimum je 5."
        ),
    )

    parser.add_argument(
        "--published-only",
        action="store_true",
        help=(
            "Vyhodnocovať iba publikované dokumenty"
        ),
    )

    parser.add_argument(
        "--max-per-document",
        type=int,
        default=1,
        help=(
            "Maximálny počet chunkov "
            "z jedného dokumentu"
        ),
    )

    parser.add_argument(
        "--strict-dataset",
        action="store_true",
        help=(
            "Ukončiť evaluáciu chybou, "
            "ak expected document nie je v indexe."
        ),
    )

    args = parser.parse_args()

    if args.limit < 5:
        parser.error(
            "--limit musí byť aspoň 5, "
            "pretože meriame Hit@5"
        )

    if args.max_per_document < 0:
        parser.error(
            "--max-per-document nesmie byť záporné"
        )

    return args


def main() -> None:
    args = parse_args()

    all_questions = load_questions(
        args.questions
    )

    split_counts = count_splits(
        all_questions
    )

    questions = (
        filter_questions_by_split(
            all_questions,
            args.split,
        )
    )

    if not questions:
        raise RuntimeError(
            f"Pre split '{args.split}' "
            "sa nenašli žiadne otázky."
        )

    print()
    print(
        "Dataset"
    )
    print(
        "=" * 60
    )

    print(
        f"Total: {len(all_questions)}"
    )

    print(
        f"Dev:   {split_counts['dev']}"
    )

    print(
        f"Test:  {split_counts['test']}"
    )

    print(
        f"Selected split: {args.split}"
    )

    print(
        f"Selected questions: {len(questions)}"
    )

    print(
        "=" * 60
    )
    print()

    indexed_documents = (
        load_index_document_paths(
            args.db
        )
    )

    missing_expected = (
        validate_dataset(
            questions,
            indexed_documents,
        )
    )

    if missing_expected:
        print()
        print(
            "POZOR: niektoré expected_documents "
            "sa nenachádzajú v indexe:"
        )

        for item in missing_expected:
            print(
                f"  {item['question_id']}: "
                f"{item['document_path']}"
            )

        print()

        if args.strict_dataset:
            raise RuntimeError(
                "Evaluačný dataset obsahuje "
                "neexistujúce expected_documents."
            )

    evaluation_rows: list[
        dict[str, Any]
    ] = []

    total = len(
        questions
    )

    for index, question in enumerate(
        questions,
        start=1,
    ):
        print(
            f"[{index:04d}/{total:04d}] "
            f"{question['id']}: "
            f"{question['question']}"
        )

        mode_results = (
            retrieve_all_modes(
                args.db,
                question[
                    "question"
                ],
                limit=args.limit,
                published_only=(
                    args.published_only
                ),
                max_per_document=(
                    args.max_per_document
                ),
            )
        )

        for mode in EVALUATION_MODES:
            evaluation_rows.append(
                evaluate_question_mode(
                    question,
                    mode,
                    mode_results[
                        mode
                    ],
                )
            )

    summary: dict[
        str,
        dict[str, Any],
    ] = {}

    by_category: dict[
        str,
        dict[
            str,
            dict[str, Any],
        ],
    ] = {}

    by_difficulty: dict[
        str,
        dict[
            str,
            dict[str, Any],
        ],
    ] = {}

    for mode in EVALUATION_MODES:
        mode_rows = [
            row
            for row in evaluation_rows
            if row[
                "mode"
            ] == mode
        ]

        summary[
            mode
        ] = aggregate_metrics(
            mode_rows
        )

        by_category[
            mode
        ] = aggregate_by_category(
            mode_rows
        )

        by_difficulty[
            mode
        ] = aggregate_by_difficulty(
            mode_rows
        )

    payload = {
        "configuration": {
            "questions_file": str(
                args.questions
            ),
            "database": str(
                args.db
            ),
            "split": (
                args.split
            ),
            "dataset_question_count": len(
                all_questions
            ),
            "selected_question_count": len(
                questions
            ),
            "dev_question_count": (
                split_counts[
                    "dev"
                ]
            ),
            "test_question_count": (
                split_counts[
                    "test"
                ]
            ),
            "limit": args.limit,
            "published_only": (
                args.published_only
            ),
            "max_per_document": (
                args.max_per_document
            ),
            "modes": list(
                EVALUATION_MODES
            ),
            "metrics": [
                "Hit@1",
                "Hit@3",
                "Hit@5",
                "MRR",
                "Recall@5",
            ],
        },
        "dataset_validation": {
            "indexed_document_count": len(
                indexed_documents
            ),
            "missing_expected_document_count": len(
                missing_expected
            ),
            "missing_expected_documents": (
                missing_expected
            ),
        },
        "summary": summary,
        "by_category": (
            by_category
        ),
        "by_difficulty": (
            by_difficulty
        ),
        "questions": (
            evaluation_rows
        ),
    }

    filename_suffix = (
        args.split
    )

    json_path = (
        args.output_dir
        / (
            "retrieval_results_"
            f"{filename_suffix}.json"
        )
    )

    csv_path = (
        args.output_dir
        / (
            "retrieval_results_"
            f"{filename_suffix}.csv"
        )
    )

    save_json_results(
        json_path,
        payload,
    )

    save_csv_results(
        csv_path,
        evaluation_rows,
    )

    print_summary(
        summary,
        split=args.split,
        selected_count=len(
            questions
        ),
        total_count=len(
            all_questions
        ),
    )

    print(
        "Výsledky:"
    )

    print(
        f"  JSON: {json_path}"
    )

    print(
        f"  CSV:  {csv_path}"
    )

    if missing_expected:
        print()
        print(
            "POZOR: pred použitím metrík "
            "v diplomovej práci oprav "
            "missing expected documents."
        )


if __name__ == "__main__":
    main()