# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Dataset loading — remote API and local CSV/JSON files."""

from __future__ import annotations

import csv
import json
import logging
from pathlib import Path
from typing import Any

import urllib.request

from klira.sdk.evals.types import KliraTestCase
from klira.sdk.utils.validation import (
    sanitize_url_path_segment,
    validate_endpoint,
    validate_path,
)

logger = logging.getLogger(__name__)


def load_dataset(
    path: str | None = None,
    *,
    dataset_id: str | None = None,
    api_key: str | None = None,
    endpoint: str | None = None,
) -> list[KliraTestCase]:
    """Load a dataset from a remote API or local file.

    Remote datasets (dataset_id) take precedence over local files (path).

    Args:
        path: Path to a local CSV or JSON file.
        dataset_id: Remote dataset ID to fetch from the API.
        api_key: API key for remote dataset fetching.
        endpoint: API endpoint base URL.

    Returns:
        List of KliraTestCase instances.

    Raises:
        ValueError: If neither path nor dataset_id is provided.
        FileNotFoundError: If local file doesn't exist.
    """
    if dataset_id is not None:
        return _load_remote(dataset_id, api_key=api_key, endpoint=endpoint)
    if path is not None:
        return _load_local(path)
    raise ValueError("Either 'path' or 'dataset_id' must be provided")


def _load_remote(
    dataset_id: str,
    *,
    api_key: str | None = None,
    endpoint: str | None = None,
) -> list[KliraTestCase]:
    """Fetch dataset from the Klira API."""
    base = endpoint or "https://api.getklira.com"
    validate_endpoint(base)
    url = f"{base}/v1/evals/datasets/{sanitize_url_path_segment(dataset_id)}"

    headers: dict[str, str] = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    req = urllib.request.Request(url, headers=headers, method="GET")
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except Exception as e:
        logger.warning("Failed to fetch remote dataset %s: %s", dataset_id, e)
        raise

    return _parse_records(data.get("test_cases", data.get("data", [])))


def _load_local(path: str) -> list[KliraTestCase]:
    """Load dataset from a local CSV or JSON file."""
    file_path = validate_path(path)
    if not file_path.exists():
        raise FileNotFoundError(f"Dataset file not found: {path}")

    suffix = file_path.suffix.lower()
    if suffix == ".csv":
        return _load_csv(file_path)
    elif suffix == ".json":
        return _load_json(file_path)
    else:
        raise ValueError(f"Unsupported file format: {suffix}. Use .csv or .json")


def _load_csv(path: Path) -> list[KliraTestCase]:
    """Load test cases from a CSV file."""
    records: list[dict[str, Any]] = []
    with open(path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            records.append(dict(row))
    return _parse_records(records)


def _load_json(path: Path) -> list[KliraTestCase]:
    """Load test cases from a JSON file."""
    with open(path, encoding="utf-8") as f:
        data = json.load(f)

    if isinstance(data, list):
        return _parse_records(data)
    if isinstance(data, dict) and "test_cases" in data:
        return _parse_records(data["test_cases"])
    raise ValueError("JSON must be a list or an object with a 'test_cases' key")


def _parse_records(records: list[dict[str, Any]]) -> list[KliraTestCase]:
    """Parse raw records into KliraTestCase instances."""
    test_cases: list[KliraTestCase] = []
    for record in records:
        input_text = record.get("input", "")
        if not input_text:
            logger.warning("Skipping test case with empty input: %s", record)
            continue

        context = record.get("context")
        if isinstance(context, str):
            context = [context]

        test_cases.append(
            KliraTestCase(
                input=input_text,
                expected_output=record.get("expected_output"),
                context=context,
                metadata=record.get("metadata"),
            )
        )
    return test_cases
