from __future__ import annotations

import argparse
import csv
import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from tools.rag.evaluation.rag_answer_data import (
    questions_fingerprint,
)
from tools.rag.evaluation.rag_metrics import (
    summarize_results,
)
from tools.rag.evaluation.rag_runner import (
    LOCAL_RAG_URL,
    OPENWEBUI_URL,
)

PARTIAL_SCHEMA_VERSION = 2

RESUME_COMPATIBILITY_KEYS = (
    "schema_version",
    "questions_file",
    "overrides_file",
    "questions_fingerprint",
    "selected_question_ids",
    "applied_override_ids",
    "split",
    "selected_question_count",
    "requested_model",
    "operation_id",
)


def _temp_path(
    path: Path,
) -> Path:
    return path.with_name(f".{path.name}." f"{os.getpid()}.tmp")


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

    temp_path = _temp_path(path)

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

            file.write("\n")
            file.flush()
            os.fsync(file.fileno())

        os.replace(
            temp_path,
            path,
        )

    finally:
        if temp_path.exists():
            temp_path.unlink()


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

    temp_path = _temp_path(path)

    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",
    ]

    try:
        with temp_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 {}

                writer.writerow(
                    {
                        "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",
                            "",
                        ),
                    }
                )

            file.flush()
            os.fsync(file.fileno())

        os.replace(
            temp_path,
            path,
        )

    finally:
        if temp_path.exists():
            temp_path.unlink()


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 build_base_result(
    *,
    question: dict[str, Any],
    question_id: str,
    question_text: str,
    selected_override_ids: list[str],
    requested_model: str,
) -> dict[str, Any]:
    return {
        "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": (requested_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"),
    }


def build_run_configuration(
    *,
    args: argparse.Namespace,
    questions: list[dict[str, Any]],
    selected_override_ids: list[str],
    operation_id: str,
) -> dict[str, Any]:
    return {
        "schema_version": (PARTIAL_SCHEMA_VERSION),
        "questions_file": str(args.questions.resolve()),
        "overrides_file": str(args.overrides.resolve()),
        "questions_fingerprint": (questions_fingerprint(questions)),
        "selected_question_ids": [str(item["id"]) for item in questions],
        "applied_override_ids": list(selected_override_ids),
        "split": args.split,
        "selected_question_count": len(questions),
        "requested_model": (args.model),
        "operation_id": (operation_id),
        "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,
    }


def load_partial_payload(
    path: Path,
) -> dict[str, Any]:
    if not path.exists():
        raise FileNotFoundError("Partial výsledok " f"neexistuje: {path}")

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

    if not isinstance(
        payload,
        dict,
    ):
        raise ValueError("Partial výsledok musí " "obsahovať JSON objekt.")

    configuration = payload.get("configuration")

    if not isinstance(
        configuration,
        dict,
    ):
        raise ValueError(
            "Partial výsledok nemá "
            "kompatibilnú configuration "
            "sekciu. Starý partial spusti "
            "znova s --overwrite."
        )

    results = payload.get("results")

    if not isinstance(
        results,
        list,
    ):
        raise ValueError("Partial výsledok nemá " "platné results pole.")

    return payload


def validate_resume_compatibility(
    partial_configuration: dict[
        str,
        Any,
    ],
    current_configuration: dict[
        str,
        Any,
    ],
) -> None:
    mismatches = [
        key
        for key in (RESUME_COMPATIBILITY_KEYS)
        if (partial_configuration.get(key) != current_configuration.get(key))
    ]

    if mismatches:
        raise ValueError(
            "Partial výsledok nie je "
            "kompatibilný s aktuálnym runom. "
            "Nezhodné polia: " + ", ".join(mismatches)
        )


def extract_resume_state(
    partial_payload: dict[str, Any],
    *,
    selected_question_ids: set[str],
) -> tuple[
    dict[str, dict[str, Any]],
    list[str],
]:
    successful: dict[
        str,
        dict[str, Any],
    ] = {}

    failed_ids: list[str] = []
    seen_ids: set[str] = set()

    for index, raw_result in enumerate(
        partial_payload["results"],
        start=1,
    ):
        if not isinstance(
            raw_result,
            dict,
        ):
            raise ValueError("Partial " f"results[{index}] " "musí byť JSON objekt.")

        question_id = raw_result.get("id")

        if (
            not isinstance(
                question_id,
                str,
            )
            or not question_id
        ):
            raise ValueError("Partial " f"results[{index}] " "nemá platné id.")

        if question_id in seen_ids:
            raise ValueError("Partial obsahuje " "duplicitné result ID: " f"{question_id}")

        seen_ids.add(question_id)

        if question_id not in selected_question_ids:
            raise ValueError("Partial obsahuje result " "mimo aktuálneho výberu: " f"{question_id}")

        if raw_result.get("error"):
            failed_ids.append(question_id)

            continue

        successful[question_id] = raw_result

    return (
        successful,
        failed_ids,
    )


def ordered_results(
    questions: list[dict[str, Any]],
    results_by_id: dict[
        str,
        dict[str, Any],
    ],
) -> list[dict[str, Any]]:
    return [
        results_by_id[str(question["id"])]
        for question in questions
        if str(question["id"]) in results_by_id
    ]


def determine_run_status(
    results: list[dict[str, Any]],
    *,
    expected_total: int,
) -> str:
    if len(results) < expected_total:
        return "partial"

    if any(result.get("error") for result in results):
        return "complete_with_errors"

    return "complete"


def build_partial_payload(
    *,
    status: str,
    configuration: dict[str, Any],
    results: list[dict[str, Any]],
    expected_total: int,
) -> dict[str, Any]:
    successful_count = sum(1 for result in results if not result.get("error"))

    return {
        "generated_at": (datetime.now(timezone.utc).isoformat()),
        "status": status,
        "configuration": (configuration),
        "attempted_so_far": len(results),
        "successful_so_far": (successful_count),
        "remaining_for_clean_completion": (expected_total - successful_count),
        "expected_total": (expected_total),
        "summary": (summarize_results(results)),
        "results": results,
    }


def prepare_output_state(
    *,
    json_path: Path,
    csv_path: Path,
    partial_path: Path,
    resume: bool,
    overwrite: bool,
) -> None:
    if resume:
        if not partial_path.exists():
            raise FileNotFoundError(
                "--resume bol zadaný, " "ale partial neexistuje: " f"{partial_path}"
            )

        return

    if overwrite:
        for path in (
            partial_path,
            json_path,
            csv_path,
        ):
            if path.exists():
                path.unlink()

        return

    if partial_path.exists():
        raise FileExistsError(
            "Existuje partial výsledok: " f"{partial_path}. " "Použi --resume alebo " "--overwrite."
        )

    existing_final = [
        path
        for path in (
            json_path,
            csv_path,
        )
        if path.exists()
    ]

    if existing_final:
        raise FileExistsError(
            "Výsledný súbor už existuje: "
            + ", ".join(str(path) for path in existing_final)
            + ". Použi --overwrite."
        )


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()