"""
cortex/eval/run_eval.py — CLI entry point for benchmark evaluation.

Usage::

    python -m cortex.eval.run_eval --dataset locomo --config cortex.toml
    python -m cortex.eval.run_eval --dataset longmemeval --data path/to/dataset.json
    python -m cortex.eval.run_eval --dataset locomo --sample  # uses bundled sample
    python -m cortex.eval.run_eval --help
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import pathlib
import sys
import time
from typing import Any, Dict, Optional


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="python -m cortex.eval.run_eval",
        description="Run Cortex memory benchmark evaluation.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python -m cortex.eval.run_eval --dataset locomo --sample
  python -m cortex.eval.run_eval --dataset locomo --config cortex.toml
  python -m cortex.eval.run_eval --dataset longmemeval --data /path/to/data.json --out results.json
  python -m cortex.eval.run_eval --dataset locomo --fuzzy-threshold 0.5 --verbose
        """,
    )
    p.add_argument(
        "--dataset",
        required=True,
        choices=["locomo", "longmemeval"],
        help="Benchmark dataset to evaluate against.",
    )
    p.add_argument(
        "--config",
        default=None,
        help="Path to cortex.toml config file (optional; uses in-memory defaults if omitted).",
    )
    p.add_argument(
        "--data",
        default=None,
        help="Path to dataset JSON file (optional; uses bundled sample if omitted).",
    )
    p.add_argument(
        "--sample",
        action="store_true",
        help="Use the bundled sample dataset (good for quick smoke-testing).",
    )
    p.add_argument(
        "--out",
        default=None,
        help="Write full results JSON to this file.",
    )
    p.add_argument(
        "--fuzzy-threshold",
        type=float,
        default=0.6,
        dest="fuzzy_threshold",
        help="Fuzzy match threshold for scoring (default: 0.6).",
    )
    p.add_argument(
        "--llm-judge",
        action="store_true",
        dest="llm_judge",
        help="Enable LLM-as-judge scoring (requires LLM config in cortex.toml).",
    )
    p.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging.",
    )
    p.add_argument(
        "--fail-under",
        type=float,
        default=None,
        dest="fail_under",
        help="Exit with code 1 if fuzzy_match_accuracy < this threshold (e.g. 0.75).",
    )
    return p


def _load_plugin(config_path: Optional[str]) -> Any:
    """Load CortexPlugin from config or create a minimal in-memory instance."""
    try:
        from cortex.api.plugin import CortexPlugin
    except ImportError as e:
        print(f"ERROR: Could not import CortexPlugin: {e}", file=sys.stderr)
        sys.exit(1)

    if config_path:
        config_file = pathlib.Path(config_path)
        if not config_file.exists():
            print(f"ERROR: Config file not found: {config_path}", file=sys.stderr)
            sys.exit(1)
        # Load from TOML config
        try:
            import tomllib  # Python 3.11+
        except ImportError:
            try:
                import tomli as tomllib  # type: ignore
            except ImportError:
                print(
                    "ERROR: tomllib not available. Install Python 3.11+ or `pip install tomli`.",
                    file=sys.stderr,
                )
                sys.exit(1)

        with open(config_file, "rb") as f:
            config_data = tomllib.load(f)

        # Extract Cortex config section
        cortex_cfg = config_data.get("cortex", config_data)

        try:
            from cortex.config import CortexConfig
            cfg = CortexConfig(**{
                k: v for k, v in cortex_cfg.items()
                if k in CortexConfig.__dataclass_fields__
            })
            plugin = CortexPlugin(config=cfg)
        except Exception as e:
            print(f"ERROR: Failed to init CortexPlugin from config: {e}", file=sys.stderr)
            sys.exit(1)
    else:
        # Minimal in-memory plugin for eval (no DB, no LLM)
        try:
            plugin = CortexPlugin()
        except Exception as e:
            print(f"ERROR: Failed to create default CortexPlugin: {e}", file=sys.stderr)
            sys.exit(1)

    return plugin


def _build_llm_judge(plugin: Any) -> Optional[Any]:
    """Build an LLM judge callable from the plugin's LLM provider."""
    try:
        llm = getattr(plugin, "llm", None)
        if llm is None:
            print("WARNING: --llm-judge requested but no LLM configured. Skipping.", file=sys.stderr)
            return None

        async def _judge_async(question: str, prediction: str, reference: str) -> float:
            prompt = (
                f"Question: {question}\n"
                f"Gold answer: {reference}\n"
                f"System answer: {prediction}\n\n"
                "Score the system answer for correctness on a scale from 0.0 to 1.0.\n"
                "0.0 = completely wrong, 1.0 = fully correct.\n"
                "Reply with ONLY a number between 0.0 and 1.0."
            )
            response = await llm.complete(prompt)
            text = str(response).strip()
            try:
                return max(0.0, min(1.0, float(text)))
            except ValueError:
                return 0.0

        # Wrap async judge in sync for harness
        def judge(question: str, prediction: str, reference: str) -> float:
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    # We're in an async context already
                    import concurrent.futures
                    with concurrent.futures.ThreadPoolExecutor() as ex:
                        future = ex.submit(asyncio.run, _judge_async(question, prediction, reference))
                        return future.result()
                return loop.run_until_complete(_judge_async(question, prediction, reference))
            except Exception:
                return 0.0

        return judge
    except Exception as e:
        print(f"WARNING: LLM judge setup failed: {e}", file=sys.stderr)
        return None


async def _run(args: argparse.Namespace) -> int:
    """Main async eval runner. Returns exit code."""
    from .harness import EvalHarness

    plugin = _load_plugin(args.config)

    llm_judge = None
    if args.llm_judge:
        llm_judge = _build_llm_judge(plugin)

    harness = EvalHarness(
        plugin=plugin,
        fuzzy_threshold=args.fuzzy_threshold,
        llm_judge=llm_judge,
    )

    # Determine data path
    data_path = None
    if args.data:
        data_path = args.data
    elif not args.sample and args.config:
        # With a real config, default to None → harness uses bundled sample
        data_path = None

    t0 = time.perf_counter()
    if args.dataset == "locomo":
        report = await harness.run_locomo(path=data_path)
    elif args.dataset == "longmemeval":
        report = await harness.run_longmemeval(path=data_path)
    else:
        print(f"ERROR: Unknown dataset: {args.dataset}", file=sys.stderr)
        return 1

    elapsed = time.perf_counter() - t0
    print(f"\nEval completed in {elapsed:.1f}s")
    print(report.summary())

    # Write full results if requested
    if args.out:
        out_path = pathlib.Path(args.out)
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(report.to_dict(), f, indent=2, default=str)
        print(f"\nResults written to: {out_path}")

    # Fail-under threshold
    if args.fail_under is not None:
        accuracy = report.metrics.get("fuzzy_match_accuracy", 0.0)
        if accuracy < args.fail_under:
            print(
                f"\nFAIL: fuzzy_match_accuracy {accuracy:.1%} < threshold {args.fail_under:.1%}",
                file=sys.stderr,
            )
            return 1
        print(f"\nPASS: fuzzy_match_accuracy {accuracy:.1%} >= threshold {args.fail_under:.1%}")

    return 0


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG, format="%(levelname)s %(name)s: %(message)s")
    else:
        logging.basicConfig(level=logging.WARNING, format="%(levelname)s: %(message)s")

    exit_code = asyncio.run(_run(args))
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
