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 evaluation.rag_answer_data import (
    questions_fingerprint,
)
from evaluation.rag_metrics import (
    summarize_results,
)
from 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()