from __future__ import annotations

import argparse
import csv
import json
import sys
import time
from datetime import datetime, timezone
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.rag_metrics import (
    evaluate_answer,
    summarize_results,
)
from evaluation.rag_runner import (
    DEFAULT_BACKOFF_BASE,
    DEFAULT_BACKOFF_MAX,
    DEFAULT_MAX_ATTEMPTS,
    DEFAULT_MODEL,
    DEFAULT_TIMEOUT,
    LOCAL_OPENAPI_URL,
    LOCAL_RAG_URL,
    OPENWEBUI_URL,
    build_rag_tool,
    load_env_value,
    request_json,
    run_question,
)
from evaluation.retrieval_runner import (
    filter_questions_by_split,
    load_questions,
)


EVALUATION_DIR = (
    PROJECT_ROOT
    / "evaluation"
)

JSON_FILES_DIR = (
    EVALUATION_DIR
    / "json_files"
)

QUESTIONS_PATH = (
    JSON_FILES_DIR
    / "questions.json"
)

RAG_ANSWER_OVERRIDES_PATH = (
    JSON_FILES_DIR
    / "rag_answer_overrides.json"
)

RESULTS_DIR = (
    EVALUATION_DIR
    / "results"
)


ALLOWED_OVERRIDE_FIELDS = {
    "question",
    "expected_answer_contains",
    "expected_source_urls",
    "should_answer",
    "note",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Answer-level RAG evaluácia "
            "cez OpenWebUI, Model120 "
            "a ZP Agent."
        )
    )

    parser.add_argument(
        "--split",
        choices=(
            "dev",
            "test",
            "all",
        ),
        default="dev",
        help=(
            "Časť datasetu. "
            "Predvolené je dev."
        ),
    )

    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help=(
            "Maximálny počet otázok "
            "na vyhodnotenie."
        ),
    )

    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=(
            "OpenWebUI model ID. "
            f"Predvolené: {DEFAULT_MODEL}."
        ),
    )

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

    parser.add_argument(
        "--overrides",
        type=Path,
        default=RAG_ANSWER_OVERRIDES_PATH,
        help=(
            "Cesta k answer-level override "
            "súboru."
        ),
    )

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

    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT,
        help=(
            "HTTP timeout v sekundách. "
            f"Predvolené: {DEFAULT_TIMEOUT}."
        ),
    )

    parser.add_argument(
        "--max-attempts",
        type=int,
        default=DEFAULT_MAX_ATTEMPTS,
        help=(
            "Maximálny počet HTTP pokusov "
            "pre retryable chyby. "
            f"Predvolené: {DEFAULT_MAX_ATTEMPTS}."
        ),
    )

    parser.add_argument(
        "--backoff-base",
        type=float,
        default=DEFAULT_BACKOFF_BASE,
        help=(
            "Počiatočný exponential backoff "
            "v sekundách. "
            f"Predvolené: {DEFAULT_BACKOFF_BASE}."
        ),
    )

    parser.add_argument(
        "--backoff-max",
        type=float,
        default=DEFAULT_BACKOFF_MAX,
        help=(
            "Maximálny exponential backoff "
            "v sekundách. "
            f"Predvolené: {DEFAULT_BACKOFF_MAX}."
        ),
    )

    parser.add_argument(
        "--delay",
        type=float,
        default=0.0,
        help=(
            "Pauza medzi otázkami "
            "v sekundách."
        ),
    )

    parser.add_argument(
        "--question-id",
        action="append",
        default=[],
        help=(
            "Vyhodnoť iba konkrétne ID. "
            "Parameter možno použiť "
            "opakovane."
        ),
    )

    parser.add_argument(
        "--fail-fast",
        action="store_true",
        help=(
            "Pri prvej chybe "
            "ukonči evaluáciu."
        ),
    )

    args = parser.parse_args()

    if (
        args.limit is not None
        and args.limit <= 0
    ):
        parser.error(
            "--limit musí byť > 0"
        )

    if args.timeout <= 0:
        parser.error(
            "--timeout musí byť > 0"
        )

    if args.max_attempts <= 0:
        parser.error(
            "--max-attempts musí byť > 0"
        )

    if args.backoff_base < 0:
        parser.error(
            "--backoff-base nesmie byť záporné"
        )

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

    if args.delay < 0:
        parser.error(
            "--delay nesmie byť záporné"
        )

    return args


def apply_question_overrides(
    questions: list[
        dict[str, Any]
    ],
    path: Path,
) -> tuple[
    list[dict[str, Any]],
    list[str],
]:
    if not path.exists():
        raise FileNotFoundError(
            "RAG answer override súbor "
            f"neexistuje: {path}"
        )

    with path.open(
        "r",
        encoding="utf-8",
    ) as file:
        overrides = json.load(
            file
        )

    if not isinstance(
        overrides,
        dict,
    ):
        raise ValueError(
            "rag_answer_overrides.json "
            "musí obsahovať JSON objekt."
        )

    known_ids = {
        str(
            item["id"]
        )
        for item in questions
    }

    unknown_ids = (
        set(
            overrides.keys()
        )
        - known_ids
    )

    if unknown_ids:
        raise ValueError(
            "Override súbor obsahuje "
            "neznáme question ID: "
            + ", ".join(
                sorted(
                    unknown_ids
                )
            )
        )

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

    applied_ids: list[
        str
    ] = []

    for question in questions:
        item = dict(
            question
        )

        question_id = str(
            item["id"]
        )

        override = overrides.get(
            question_id
        )

        if override is None:
            result.append(
                item
            )

            continue

        if not isinstance(
            override,
            dict,
        ):
            raise ValueError(
                f"{question_id}: "
                "override musí byť "
                "JSON objekt."
            )

        unsupported_fields = (
            set(
                override.keys()
            )
            - ALLOWED_OVERRIDE_FIELDS
        )

        if unsupported_fields:
            raise ValueError(
                f"{question_id}: "
                "nepovolené override polia: "
                + ", ".join(
                    sorted(
                        unsupported_fields
                    )
                )
            )

        if "question" in override:
            overridden_question = (
                override[
                    "question"
                ]
            )

            if (
                not isinstance(
                    overridden_question,
                    str,
                )
                or not overridden_question.strip()
            ):
                raise ValueError(
                    f"{question_id}: "
                    "override question musí byť "
                    "neprázdny reťazec."
                )

        item.update(
            override
        )

        result.append(
            item
        )

        applied_ids.append(
            question_id
        )

    return (
        result,
        applied_ids,
    )


def select_questions(
    questions: list[
        dict[str, Any]
    ],
    *,
    split: str,
    question_ids: list[str],
    limit: int | None,
) -> list[
    dict[str, Any]
]:
    selected = (
        filter_questions_by_split(
            questions,
            split,
        )
    )

    if question_ids:
        wanted = set(
            question_ids
        )

        selected = [
            item
            for item in selected
            if item.get(
                "id"
            ) in wanted
        ]

    if limit is not None:
        selected = selected[
            :limit
        ]

    return selected


def result_prefix(
    *,
    split: str,
    limit: int | None,
    question_ids: list[str],
) -> str:
    prefix = (
        f"rag_answers_{split}"
    )

    if question_ids:
        return (
            prefix
            + "_selected"
        )

    if limit is not None:
        return (
            prefix
            + f"_limit{limit}"
        )

    return prefix


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,
    results: list[
        dict[str, Any]
    ],
) -> None:
    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    fieldnames = [
        "id",
        "split",
        "category",
        "difficulty",
        "question",
        "requested_model",
        "response_model",
        "tool_called",
        "tool_call_count",
        "answer_contains_score",
        "source_url_score",
        "should_answer_score",
        "overall_score",
        "strict_pass",
        "expected_answer_contains",
        "expected_source_urls",
        "rag_source_urls",
        "answer",
        "first_model_latency_seconds",
        "tool_latency_seconds",
        "final_model_latency_seconds",
        "total_latency_seconds",
        "prompt_tokens",
        "completion_tokens",
        "total_tokens",
        "error",
    ]

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

        writer.writeheader()

        for result in results:
            usage = (
                result.get(
                    "usage"
                )
                or {}
            )

            row = {
                "id": result.get(
                    "id"
                ),
                "split": result.get(
                    "split"
                ),
                "category": result.get(
                    "category"
                ),
                "difficulty": result.get(
                    "difficulty"
                ),
                "question": result.get(
                    "question"
                ),
                "requested_model": result.get(
                    "requested_model"
                ),
                "response_model": result.get(
                    "response_model"
                ),
                "tool_called": result.get(
                    "tool_called"
                ),
                "tool_call_count": result.get(
                    "tool_call_count"
                ),
                "answer_contains_score": result.get(
                    "answer_contains_score"
                ),
                "source_url_score": result.get(
                    "source_url_score"
                ),
                "should_answer_score": result.get(
                    "should_answer_score"
                ),
                "overall_score": result.get(
                    "overall_score"
                ),
                "strict_pass": result.get(
                    "strict_pass"
                ),
                "expected_answer_contains": (
                    json.dumps(
                        result.get(
                            "expected_answer_contains",
                            [],
                        ),
                        ensure_ascii=False,
                    )
                ),
                "expected_source_urls": (
                    json.dumps(
                        result.get(
                            "expected_source_urls",
                            [],
                        ),
                        ensure_ascii=False,
                    )
                ),
                "rag_source_urls": (
                    json.dumps(
                        result.get(
                            "rag_source_urls",
                            [],
                        ),
                        ensure_ascii=False,
                    )
                ),
                "answer": result.get(
                    "answer",
                    "",
                ),
                "first_model_latency_seconds": (
                    result.get(
                        "first_model_latency_seconds"
                    )
                ),
                "tool_latency_seconds": result.get(
                    "tool_latency_seconds"
                ),
                "final_model_latency_seconds": (
                    result.get(
                        "final_model_latency_seconds"
                    )
                ),
                "total_latency_seconds": result.get(
                    "total_latency_seconds"
                ),
                "prompt_tokens": usage.get(
                    "prompt_tokens",
                    0,
                ),
                "completion_tokens": usage.get(
                    "completion_tokens",
                    0,
                ),
                "total_tokens": usage.get(
                    "total_tokens",
                    0,
                ),
                "error": result.get(
                    "error",
                    "",
                ),
            }

            writer.writerow(
                row
            )


def build_error_result(
    base_result: dict[str, Any],
    exc: Exception,
) -> dict[str, Any]:
    return {
        **base_result,
        "answer": "",
        "tool_called": False,
        "tool_call_count": 0,
        "tool_calls": [],
        "rag_source_urls": [],
        "first_model_latency_seconds": 0.0,
        "tool_latency_seconds": 0.0,
        "final_model_latency_seconds": 0.0,
        "total_latency_seconds": 0.0,
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        },
        "response_model": None,
        "answer_matches": [],
        "answer_contains_score": 0.0,
        "source_matches": [],
        "source_url_score": 0.0,
        "should_answer_ok": False,
        "should_answer_score": 0.0,
        "returned_no_answer": False,
        "tool_score": 0.0,
        "overall_score": 0.0,
        "strict_pass": False,
        "error": (
            f"{type(exc).__name__}: "
            f"{exc}"
        ),
    }


def print_summary(
    summary: dict[str, Any],
) -> None:
    print()
    print(
        "RAG answer evaluation"
    )

    print(
        "=" * 78
    )

    print(
        f"Total: "
        f"{summary['total']}"
    )

    print(
        f"Completed: "
        f"{summary['completed']}"
    )

    print(
        f"Errors: "
        f"{summary['errors']}"
    )

    print(
        "Tool call rate: "
        f"{summary['tool_call_rate']:.3f}"
    )

    print(
        "Answer contains: "
        f"{summary['answer_contains_score']:.3f}"
    )

    print(
        "Source URL: "
        f"{summary['source_url_score']:.3f}"
    )

    print(
        "Should answer: "
        f"{summary['should_answer_score']:.3f}"
    )

    print(
        "Overall: "
        f"{summary['overall_score']:.3f}"
    )

    print(
        "Strict pass: "
        f"{summary['strict_pass_count']}"
        f"/{summary['total']} "
        f"({summary['strict_pass_rate']:.3f})"
    )

    print(
        "Mean latency: "
        f"{summary['mean_latency_seconds']:.3f} s"
    )

    print(
        "Total tokens: "
        f"{summary['total_tokens']}"
    )

    print(
        "=" * 78
    )

    print()


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

    try:
        all_questions = (
            load_questions(
                args.questions
            )
        )

        (
            all_questions,
            applied_override_ids,
        ) = apply_question_overrides(
            all_questions,
            args.overrides,
        )

    except (
        FileNotFoundError,
        ValueError,
        json.JSONDecodeError,
    ) as exc:
        print(
            f"ERROR: {exc}",
            file=sys.stderr,
        )

        return 2

    questions = select_questions(
        all_questions,
        split=args.split,
        question_ids=(
            args.question_id
        ),
        limit=args.limit,
    )

    if not questions:
        print(
            "ERROR: Žiadne otázky "
            "nezodpovedajú filtru.",
            file=sys.stderr,
        )

        return 2

    selected_ids = {
        str(
            question["id"]
        )
        for question in questions
    }

    selected_override_ids = [
        question_id
        for question_id in applied_override_ids
        if question_id in selected_ids
    ]

    try:
        openwebui_api_key = (
            load_env_value(
                "OPENWEBUI_API_KEY",
                allow_env_file=False,
            )
        )

        search_api_key = (
            load_env_value(
                "SEARCH_API_KEY"
            )
        )

    except RuntimeError as exc:
        print(
            f"ERROR: {exc}",
            file=sys.stderr,
        )

        return 2

    try:
        openapi = request_json(
            LOCAL_OPENAPI_URL,
            timeout=args.timeout,
            max_attempts=(
                args.max_attempts
            ),
            backoff_base=(
                args.backoff_base
            ),
            backoff_max=(
                args.backoff_max
            ),
        )

        operation_id, rag_tool = (
            build_rag_tool(
                openapi
            )
        )

    except RuntimeError as exc:
        print(
            f"ERROR: {exc}",
            file=sys.stderr,
        )

        return 2

    prefix = result_prefix(
        split=args.split,
        limit=args.limit,
        question_ids=(
            args.question_id
        ),
    )

    json_path = (
        args.results_dir
        / f"{prefix}.json"
    )

    csv_path = (
        args.results_dir
        / f"{prefix}.csv"
    )

    partial_path = (
        args.results_dir
        / f"{prefix}.partial.json"
    )

    print()
    print(
        "RAG answer evaluation"
    )

    print(
        "=" * 60
    )

    print(
        f"Dataset: "
        f"{args.questions}"
    )

    print(
        f"Overrides: "
        f"{args.overrides}"
    )

    print(
        "Applied overrides: "
        f"{len(selected_override_ids)}"
    )

    if selected_override_ids:
        print(
            "Override IDs: "
            + ", ".join(
                selected_override_ids
            )
        )

    print(
        f"Model: "
        f"{args.model}"
    )

    print(
        f"Tool: "
        f"{operation_id}"
    )

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

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

    print(
        "Retry: "
        f"{args.max_attempts} pokusy"
    )

    print(
        "Backoff: "
        f"{args.backoff_base}s "
        f"→ max {args.backoff_max}s"
    )

    print(
        "=" * 60
    )

    print()

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

    total = len(
        questions
    )

    stopped_early = False

    for index, question in enumerate(
        questions,
        start=1,
    ):
        question_id = str(
            question[
                "id"
            ]
        )

        question_text = str(
            question[
                "question"
            ]
        )

        print(
            f"[{index:04d}/{total:04d}] "
            f"{question_id}: "
            f"{question_text}"
        )

        base_result: dict[
            str,
            Any,
        ] = {
            "id": question_id,
            "split": question.get(
                "split"
            ),
            "category": question.get(
                "category",
                "unknown",
            ),
            "difficulty": question.get(
                "difficulty",
                "unknown",
            ),
            "question": question_text,
            "override_applied": (
                question_id
                in selected_override_ids
            ),
            "requested_model": (
                args.model
            ),
            "expected_documents": question.get(
                "expected_documents",
                [],
            ),
            "expected_source_urls": question.get(
                "expected_source_urls",
                [],
            ),
            "expected_answer_contains": question.get(
                "expected_answer_contains",
                [],
            ),
            "should_answer": question.get(
                "should_answer",
                True,
            ),
            "note": question.get(
                "note"
            ),
        }

        try:
            run_result = run_question(
                question_text,
                model=args.model,
                operation_id=(
                    operation_id
                ),
                rag_tool=rag_tool,
                openwebui_api_key=(
                    openwebui_api_key
                ),
                search_api_key=(
                    search_api_key
                ),
                timeout=args.timeout,
                max_attempts=(
                    args.max_attempts
                ),
                backoff_base=(
                    args.backoff_base
                ),
                backoff_max=(
                    args.backoff_max
                ),
            )

            scores = evaluate_answer(
                question,
                run_result[
                    "answer"
                ],
                tool_called=bool(
                    run_result[
                        "tool_called"
                    ]
                ),
            )

            result = {
                **base_result,
                **run_result,
                **scores,
                "error": None,
            }

            print(
                "    tool:",
                result[
                    "tool_called"
                ],
                "| answer:",
                f"{result['answer_contains_score']:.2f}",
                "| source:",
                f"{result['source_url_score']:.2f}",
                "| strict:",
                result[
                    "strict_pass"
                ],
                "| latency:",
                f"{result['total_latency_seconds']:.2f}s",
            )

        except Exception as exc:
            result = (
                build_error_result(
                    base_result,
                    exc,
                )
            )

            print(
                "    ERROR:",
                result[
                    "error"
                ],
            )

            if args.fail_fast:
                stopped_early = True

        results.append(
            result
        )

        partial_payload = {
            "generated_at": (
                datetime.now(
                    timezone.utc
                ).isoformat()
            ),
            "status": "partial",
            "split": (
                args.split
            ),
            "requested_model": (
                args.model
            ),
            "questions_file": str(
                args.questions
            ),
            "overrides_file": str(
                args.overrides
            ),
            "applied_override_ids": (
                selected_override_ids
            ),
            "max_attempts": (
                args.max_attempts
            ),
            "backoff_base": (
                args.backoff_base
            ),
            "backoff_max": (
                args.backoff_max
            ),
            "completed_so_far": len(
                results
            ),
            "expected_total": (
                total
            ),
            "summary": (
                summarize_results(
                    results
                )
            ),
            "results": (
                results
            ),
        }

        save_json_results(
            partial_path,
            partial_payload,
        )

        if stopped_early:
            break

        if (
            args.delay > 0
            and index < total
        ):
            time.sleep(
                args.delay
            )

    summary = (
        summarize_results(
            results
        )
    )

    completed_all = (
        len(results)
        == total
    )

    payload = {
        "generated_at": (
            datetime.now(
                timezone.utc
            ).isoformat()
        ),
        "status": (
            "complete"
            if completed_all
            else "partial"
        ),
        "configuration": {
            "questions_file": str(
                args.questions
            ),
            "overrides_file": str(
                args.overrides
            ),
            "applied_override_ids": (
                selected_override_ids
            ),
            "split": (
                args.split
            ),
            "selected_question_count": (
                total
            ),
            "requested_model": (
                args.model
            ),
            "openwebui_url": (
                OPENWEBUI_URL
            ),
            "rag_url": (
                LOCAL_RAG_URL
            ),
            "timeout": (
                args.timeout
            ),
            "max_attempts": (
                args.max_attempts
            ),
            "backoff_base": (
                args.backoff_base
            ),
            "backoff_max": (
                args.backoff_max
            ),
            "delay": (
                args.delay
            ),
        },
        "summary": (
            summary
        ),
        "results": (
            results
        ),
    }

    save_json_results(
        json_path,
        payload,
    )

    save_csv_results(
        csv_path,
        results,
    )

    if (
        partial_path.exists()
        and completed_all
    ):
        partial_path.unlink()

    print_summary(
        summary
    )

    print(
        "Výsledky:"
    )

    print(
        f"  JSON: "
        f"{json_path}"
    )

    print(
        f"  CSV:  "
        f"{csv_path}"
    )

    if summary[
        "errors"
    ]:
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(
        main()
    )