from __future__ import annotations

import argparse
import json
import os
import re
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

from env.environment import LastMileDeliveryEnvironment
from env.models import Action, ActionType, Direction, Observation, Order, OrderPriority, OrderStatus, SimulatorAction, StepInfo, StepResult
from env.utils import in_bounds, manhattan, move_coordinate
from grader.grader import DeliveryEpisodeGrader
from tasks.registry import get_task_config, get_task_definition, list_tasks

try:
	from openai import OpenAI
except Exception:  # pragma: no cover - import guard for environments without dependency installed
	OpenAI = None  # type: ignore[assignment]


DEFAULT_SEED = 42
MAX_RUNTIME_SECONDS = 19 * 60
DEFAULT_TASK = "easy"
DEFAULT_MODEL_NAME = "gpt-4o-mini"

DIRECTION_PRIORITY: List[Direction] = [
	Direction.UP,
	Direction.LEFT,
	Direction.RIGHT,
	Direction.DOWN,
]


@dataclass(frozen=True)
class RuntimeSettings:
	openai_api_key: str
	model_name: str
	api_base_url: Optional[str]
	use_remote_model: bool


def _read_env_first(*names: str) -> str:
	for name in names:
		value = (os.getenv(name) or "").strip()
		if value:
			return value
	return ""


def _require_any_env(*names: str) -> str:
	value = _read_env_first(*names)
	if value:
		return value
	raise RuntimeError(f"Missing required environment variable. Set one of: {', '.join(names)}")


def _read_bool_env(name: str, default: bool = False) -> bool:
	raw = (os.getenv(name) or "").strip().lower()
	if not raw:
		return default
	return raw in {"1", "true", "yes", "on"}


def _resolve_model_name() -> str:
	value = _read_env_first("MODEL_NAME", "OPENAI_MODEL")
	return value or DEFAULT_MODEL_NAME


def load_runtime_settings() -> RuntimeSettings:
	base_url = _read_env_first("API_BASE_URL", "OPENAI_BASE_URL") or None
	model_name = _resolve_model_name()

	return RuntimeSettings(
		openai_api_key=_require_any_env("HF_TOKEN", "OPENAI_API_KEY"),
		model_name=model_name,
		api_base_url=base_url,
		use_remote_model=_read_bool_env("OPENAI_USE_MODEL", default=False),
	)


class DeterministicHeuristicPolicy:
	def __init__(self, settings: RuntimeSettings, seed: int):
		if OpenAI is None:
			raise RuntimeError("openai package is required but not installed")

		self.model_name = settings.model_name
		self.use_remote_model = settings.use_remote_model
		self.seed = seed
		self._recent_positions: List[Tuple[int, int]] = []
		self._previous_position: Optional[Tuple[int, int]] = None
		# Keep a deterministic fallback path so local validation does not depend on external API availability.
		self.client: Any | None = None
		self.client_init_error: str | None = None
		try:
			client_kwargs: Dict[str, Any] = {
				"api_key": settings.openai_api_key,
				"timeout": 20.0,
			}
			if settings.api_base_url is not None:
				client_kwargs["base_url"] = settings.api_base_url
			self.client = OpenAI(**client_kwargs)
		except Exception as exc:
			self.client_init_error = type(exc).__name__

	def _update_navigation_memory(self, observation: Observation, step_index: int) -> None:
		current = (observation.agent_location.x, observation.agent_location.y)

		# Step index resets to 1 at each new episode; clear memory so repeated single-step
		# inference calls remain deterministic.
		if step_index <= 1:
			self._recent_positions = [current]
			self._previous_position = None
			return

		self._recent_positions.append(current)
		if len(self._recent_positions) > 8:
			self._recent_positions.pop(0)

		if len(self._recent_positions) >= 2:
			self._previous_position = self._recent_positions[-2]
		else:
			self._previous_position = None

	def _is_local_oscillation(self) -> bool:
		if len(self._recent_positions) < 4:
			return False

		a, b, c, d = self._recent_positions[-4:]
		return a == c and b == d and c != d

	def decide(self, observation: Observation, step_index: int) -> Tuple[SimulatorAction, Optional[str]]:
		self._update_navigation_memory(observation=observation, step_index=step_index)
		break_oscillation = self._is_local_oscillation()

		if self.use_remote_model and self.client is not None:
			model_action, model_error = self._try_model_action(observation, step_index=step_index)
			if model_action is not None:
				return model_action, model_error

		try:
			action = _heuristic_action(
				observation,
				previous_position=self._previous_position,
				break_oscillation=break_oscillation,
			)
			return action, self.client_init_error if self.use_remote_model and self.client is None else None
		except Exception as exc:
			fallback = SimulatorAction(action_type=ActionType.WAIT)
			return fallback, f"policy_error:{type(exc).__name__}"

	def _try_model_action(self, observation: Observation, step_index: int) -> Tuple[Optional[SimulatorAction], Optional[str]]:
		if self.client is None:
			return None, "openai_client_unavailable"

		system_prompt = (
			"You are a deterministic last-mile dispatch policy. "
			"Return one JSON object matching this schema exactly: "
			"{\"move\": string|null, \"accept_order\": integer|null, \"deliver_order\": boolean, \"wait\": boolean}. "
			"Exactly one intent must be active. Do not include markdown."
		)
		payload = {
			"step_index": step_index,
			"observation": observation.model_dump(),
		}

		try:
			response = self.client.responses.create(
				model=self.model_name,
				temperature=0,
				max_output_tokens=120,
				input=[
					{
						"role": "system",
						"content": [{"type": "text", "text": system_prompt}],
					},
					{
						"role": "user",
						"content": [{"type": "text", "text": json.dumps(payload, separators=(",", ":"), sort_keys=True)}],
					},
				],
			)
		except Exception as exc:
			return None, f"model_request_error:{type(exc).__name__}"

		output_text = getattr(response, "output_text", "") or ""
		parsed_json = self._extract_json_object(output_text)
		if parsed_json is None:
			return None, "model_parse_error"

		try:
			action = Action.model_validate(parsed_json).to_simulator_action()
		except Exception as exc:
			return None, f"model_action_invalid:{type(exc).__name__}"

		return action, None

	@staticmethod
	def _extract_json_object(text: str) -> Optional[Dict[str, Any]]:
		candidate = text.strip()
		if not candidate:
			return None

		if candidate.startswith("```"):
			candidate = re.sub(r"^```[a-zA-Z0-9_\-]*\s*", "", candidate)
			candidate = re.sub(r"\s*```$", "", candidate).strip()

		try:
			value = json.loads(candidate)
			if isinstance(value, dict):
				return value
		except Exception:
			pass

		start = candidate.find("{")
		end = candidate.rfind("}")
		if start == -1 or end == -1 or end <= start:
			return None

		try:
			value = json.loads(candidate[start : end + 1])
			if isinstance(value, dict):
				return value
		except Exception:
			return None

		return None


_POLICY: DeterministicHeuristicPolicy | None = None


def _get_policy(seed: int = DEFAULT_SEED) -> DeterministicHeuristicPolicy:
	global _POLICY
	if _POLICY is None:
		settings = load_runtime_settings()
		_POLICY = DeterministicHeuristicPolicy(settings=settings, seed=seed)
	return _POLICY


def _priority_rank(order: Order) -> int:
	return 0 if order.priority == OrderPriority.HIGH else 1


def _choose_pending_order(observation: Observation) -> Optional[Order]:
	agent = (observation.agent_location.x, observation.agent_location.y)
	pending = [order for order in observation.pending_orders if order.status == OrderStatus.PENDING]
	if not pending:
		return None

	pending.sort(
		key=lambda order: (
			_priority_rank(order),
			manhattan(agent, (order.pickup.x, order.pickup.y)),
			order.order_id,
		)
	)
	return pending[0]


def _current_target(observation: Observation) -> Optional[Tuple[int, int]]:
	current_order = observation.current_order
	if current_order is None:
		return None

	agent = (observation.agent_location.x, observation.agent_location.y)
	if current_order.status == OrderStatus.ACCEPTED:
		return (current_order.pickup.x, current_order.pickup.y)

	if current_order.status == OrderStatus.PICKED_UP:
		targets = current_order.delivery_locations or [current_order.dropoff]
		target_coords = [(target.x, target.y) for target in targets]
		target_coords.sort(key=lambda coord: (manhattan(agent, coord), coord[0], coord[1]))
		return target_coords[0]

	if current_order.status == OrderStatus.PENDING:
		return (current_order.pickup.x, current_order.pickup.y)

	return None


def _nearest_charging_station(observation: Observation) -> Optional[Tuple[int, int]]:
	if not observation.charging_stations:
		return None

	agent = (observation.agent_location.x, observation.agent_location.y)
	station_coords = [(item.x, item.y) for item in observation.charging_stations]
	station_coords.sort(key=lambda coord: (manhattan(agent, coord), coord[0], coord[1]))
	return station_coords[0]


def _best_move_direction(
	observation: Observation,
	target: Tuple[int, int],
	previous_position: Optional[Tuple[int, int]] = None,
	break_oscillation: bool = False,
) -> Optional[Direction]:
	current = (observation.agent_location.x, observation.agent_location.y)
	blocked = {(item.x, item.y) for item in observation.obstacles}
	blocked.update((item.x, item.y) for item in observation.dynamic_obstacles)

	candidates: List[Tuple[int, int, Direction, Tuple[int, int]]] = []
	for index, direction in enumerate(DIRECTION_PRIORITY):
		nxt = move_coordinate(current, direction)
		if not in_bounds(nxt, observation.grid_width, observation.grid_height):
			continue
		if nxt in blocked:
			continue

		distance = manhattan(nxt, target)
		candidates.append((distance, index, direction, nxt))

	if not candidates:
		return None

	candidates.sort(key=lambda item: (item[0], item[1]))

	if previous_position is not None and len(candidates) > 1:
		best_distance = candidates[0][0]
		for distance, _, direction, nxt in candidates:
			if nxt == previous_position:
				continue

			# Stronger anti-backtrack when we detect A-B-A-B style oscillation.
			if break_oscillation:
				return direction

			# Prefer a non-backtracking move when it is near-optimal.
			if distance <= best_distance + 1:
				return direction

			break

	return candidates[0][2]


def _move_toward_target(
	observation: Observation,
	target: Tuple[int, int],
	previous_position: Optional[Tuple[int, int]] = None,
	break_oscillation: bool = False,
) -> SimulatorAction:
	direction = _best_move_direction(
		observation,
		target,
		previous_position=previous_position,
		break_oscillation=break_oscillation,
	)
	if direction is None:
		return SimulatorAction(action_type=ActionType.WAIT)
	return SimulatorAction(action_type=ActionType.MOVE, direction=direction)


def _battery_aware_action(
	observation: Observation,
	target: Tuple[int, int],
	previous_position: Optional[Tuple[int, int]] = None,
	break_oscillation: bool = False,
) -> Optional[SimulatorAction]:
	battery_level = observation.battery_level
	if battery_level is None:
		return None

	# Keep pickup behavior aggressive; only interrupt for charging when battery is near-empty.
	if observation.current_order is None and battery_level > 3:
		return None

	charger_target = _nearest_charging_station(observation)
	if charger_target is None:
		return None

	agent = (observation.agent_location.x, observation.agent_location.y)
	distance_to_charger = manhattan(agent, charger_target)
	distance_to_target = manhattan(agent, target)

	# Trigger recharge only when battery is close to the minimum needed to reach a charger.
	recharge_trigger = distance_to_charger + 1
	if battery_level > recharge_trigger:
		return None

	if agent == charger_target:
		# Recharge just enough to leave and make forward progress; avoid long waiting loops.
		resume_level = min(12, max(8, distance_to_target + 2))
		if battery_level < resume_level:
			return SimulatorAction(action_type=ActionType.WAIT)
		return None

	return _move_toward_target(
		observation,
		charger_target,
		previous_position=previous_position,
		break_oscillation=break_oscillation,
	)


def _heuristic_action(
	observation: Observation,
	previous_position: Optional[Tuple[int, int]] = None,
	break_oscillation: bool = False,
) -> SimulatorAction:
	agent = (observation.agent_location.x, observation.agent_location.y)
	current_order = observation.current_order

	if current_order is None:
		selected = _choose_pending_order(observation)
		if selected is None:
			return SimulatorAction(action_type=ActionType.WAIT)

		pickup_target = (selected.pickup.x, selected.pickup.y)
		if agent == pickup_target:
			return SimulatorAction(action_type=ActionType.ACCEPT_ORDER, order_id=selected.order_id)

		recharge_action = _battery_aware_action(
			observation,
			pickup_target,
			previous_position=previous_position,
			break_oscillation=break_oscillation,
		)
		if recharge_action is not None:
			return recharge_action

		# Move toward the nearest pickup instead of idling while no order is selected.
		return _move_toward_target(
			observation,
			pickup_target,
			previous_position=previous_position,
			break_oscillation=break_oscillation,
		)

	target = _current_target(observation)
	if target is None:
		selected = _choose_pending_order(observation)
		if selected is None:
			return SimulatorAction(action_type=ActionType.WAIT)
		return _move_toward_target(
			observation,
			(selected.pickup.x, selected.pickup.y),
			previous_position=previous_position,
			break_oscillation=break_oscillation,
		)

	if agent == target:
		return SimulatorAction(action_type=ActionType.DELIVER_ORDER)

	recharge_action = _battery_aware_action(
		observation,
		target,
		previous_position=previous_position,
		break_oscillation=break_oscillation,
	)
	if recharge_action is not None:
		return recharge_action

	# With an active target, prefer movement over waiting whenever a valid move exists.
	return _move_toward_target(
		observation,
		target,
		previous_position=previous_position,
		break_oscillation=break_oscillation,
	)


def _bool_literal(value: bool) -> str:
	return "true" if value else "false"


def _format_error(error: Optional[str]) -> str:
	if error is None:
		return "null"
	return error.replace(" ", "_")


def _format_token(value: Optional[str], fallback: str) -> str:
	token = (value or "").strip()
	if not token:
		token = fallback
	return "_".join(token.split())


def _format_action(action: SimulatorAction) -> str:
	payload = Action.from_simulator_action(action).model_dump()
	return json.dumps(payload, separators=(",", ":"), sort_keys=True)


def format_start_line(task: str, env_name: str, model_name: str) -> str:
	return (
		f"[START] "
		f"task={_format_token(task, fallback='unknown_task')} "
		f"env={_format_token(env_name, fallback='unknown_env')} "
		f"model={_format_token(model_name, fallback='unknown_model')}"
	)


def format_step_line(step_index: int, action: SimulatorAction, reward: float, done: bool, error: Optional[str]) -> str:
	return (
		f"[STEP] step={step_index} "
		f"action={_format_action(action)} "
		f"reward={reward:.2f} "
		f"done={_bool_literal(done)} "
		f"error={_format_error(error)}"
	)


def _format_reward_series(rewards: Sequence[float] | float) -> str:
	if isinstance(rewards, (int, float)):
		return f"{float(rewards):.2f}"

	reward_values = list(rewards)
	if not reward_values:
		return "0.00"
	return ",".join(f"{value:.2f}" for value in reward_values)


def format_end_line(success: bool, steps: int, rewards: Sequence[float] | float) -> str:
	return f"[END] success={_bool_literal(success)} steps={steps} rewards={_format_reward_series(rewards)}"


def _episode_success(done: bool, observation: Observation) -> bool:
	if not done:
		return False
	return len(observation.pending_orders) == 0 and observation.current_order is None


def _grade_score(task: str, trajectory: Sequence[StepResult], final_observation: Observation) -> float:
	grader = DeliveryEpisodeGrader()
	success_condition: Dict[str, Any] = {}
	try:
		task_definition = get_task_definition(task)
		success_condition = dict(task_definition.get("success_condition", {}))
	except Exception:
		success_condition = {}

	report = grader.grade_episode(
		trajectory=trajectory,
		final_observation=final_observation,
		success_condition=success_condition,
	)
	return float(report.score)


def run_episode(
	task: str = DEFAULT_TASK,
	seed: int = DEFAULT_SEED,
	max_steps: Optional[int] = None,
	trace: bool = True,
) -> Dict[str, Any]:
	steps = 0
	reward_history: List[float] = []
	trajectory: List[StepResult] = []
	success = False
	step_printed = False
	startup_error: Optional[str] = None
	runtime_error: Optional[str] = None
	model_name = _resolve_model_name()
	env_name = LastMileDeliveryEnvironment.__name__
	settings: Optional[RuntimeSettings] = None
	score = 0.0
	env: Optional[LastMileDeliveryEnvironment] = None
	observation: Optional[Observation] = None

	try:
		try:
			settings = load_runtime_settings()
			model_name = settings.model_name
		except Exception as exc:
			startup_error = f"startup_error:{type(exc).__name__}"

		if trace:
			print(format_start_line(task=task, env_name=env_name, model_name=model_name))

		if settings is None:
			action = SimulatorAction(action_type=ActionType.WAIT)
			if trace:
				print(
					format_step_line(
						step_index=1,
						action=action,
						reward=0.0,
						done=False,
						error=startup_error,
					)
				)
			step_printed = True
			return {
				"success": False,
				"steps": 0,
				"rewards": [],
				"total_reward": 0.0,
				"score": 0.0,
			}

		policy = DeterministicHeuristicPolicy(settings=settings, seed=seed)
		config = get_task_config(task_name=task, seed=seed)
		env = LastMileDeliveryEnvironment(config=config)
		observation = env.reset()

		start_time = time.monotonic()
		configured_limit = max_steps if max_steps is not None else config.max_steps
		step_limit = max(1, min(configured_limit, config.max_steps))

		for step_index in range(1, step_limit + 1):
			if time.monotonic() - start_time >= MAX_RUNTIME_SECONDS:
				runtime_error = "timeout"
				break

			action, action_error = policy.decide(observation, step_index=step_index)
			try:
				step_result = env.step_result(action)
				trajectory.append(step_result)
				observation = step_result.observation
				reward = float(step_result.reward)
				done = step_result.done
				info = step_result.info.model_dump()
			except Exception as exc:
				reward = 0.0
				done = False
				info = {"invalid_action": True, "message": f"step_error:{type(exc).__name__}"}
				action = SimulatorAction(action_type=ActionType.WAIT)
				runtime_error = str(info["message"])
				if env is not None:
					fallback_observation = env.current_observation()
					trajectory.append(
						StepResult(
							observation=fallback_observation,
							reward=reward,
							done=done,
							info=StepInfo(invalid_action=True, message=str(info["message"])),
						)
					)

			step_error = action_error
			if step_error is None and bool(info.get("invalid_action")):
				step_error = str(info.get("message", "invalid_action"))

			steps = step_index
			reward_history.append(float(reward))
			if trace:
				print(
					format_step_line(
						step_index=step_index,
						action=action,
						reward=reward,
						done=done,
						error=step_error,
					)
				)
			step_printed = True

			if done:
				success = _episode_success(done=done, observation=observation)
				break
			if runtime_error is not None:
				break

		if env is not None:
			final_observation = env.current_observation()
			score = _grade_score(task=task, trajectory=trajectory, final_observation=final_observation)

		return {
			"success": success,
			"steps": steps,
			"rewards": [round(value, 2) for value in reward_history],
			"total_reward": round(sum(reward_history), 2),
			"score": round(score, 4),
		}

	except BaseException as exc:
		runtime_error = f"fatal_error:{type(exc).__name__}"
		return {
			"success": False,
			"steps": steps,
			"rewards": [round(value, 2) for value in reward_history],
			"total_reward": round(sum(reward_history), 2),
			"score": 0.0,
		}
	finally:
		if trace and not step_printed:
			action = SimulatorAction(action_type=ActionType.WAIT)
			fallback_error = runtime_error or startup_error
			print(
				format_step_line(
					step_index=max(1, steps + 1),
					action=action,
					reward=0.0,
					done=False,
					error=fallback_error,
				)
			)

		if trace:
			print(format_end_line(success=success, steps=steps, rewards=reward_history))


def run_baseline_suite(seed: int = DEFAULT_SEED, max_steps: Optional[int] = None) -> Dict[str, Any]:
	task_results: List[Dict[str, Any]] = []
	for task_name in list_tasks():
		episode_result = run_episode(task=task_name, seed=seed, max_steps=max_steps, trace=False)
		task_results.append(
			{
				"task": task_name,
				"success": bool(episode_result.get("success", False)),
				"steps": int(episode_result.get("steps", 0)),
				"total_reward": float(episode_result.get("total_reward", 0.0)),
				"score": float(episode_result.get("score", 0.0)),
			}
		)

	average_score = 0.0
	if task_results:
		average_score = round(sum(item["score"] for item in task_results) / len(task_results), 4)

	return {
		"seed": seed,
		"model": _resolve_model_name(),
		"tasks": task_results,
		"average_score": average_score,
	}


def inference(observation: Dict[str, Any]) -> Dict[str, Any]:
	"""OpenEnv action entrypoint with deterministic heuristic policy and OpenAI client setup."""
	parsed = Observation.model_validate(observation)
	try:
		action, _ = _get_policy(seed=DEFAULT_SEED).decide(parsed, step_index=1)
	except Exception:
		action = _heuristic_action(parsed)
	return Action.from_simulator_action(action).model_dump()


def predict(observation: Dict[str, Any]) -> Dict[str, Any]:
	"""Compatibility alias for runners expecting predict(...)."""
	return inference(observation)


def _build_arg_parser() -> argparse.ArgumentParser:
	parser = argparse.ArgumentParser(description="Run deterministic OpenEnv inference episode")
	parser.add_argument("--task", default=DEFAULT_TASK, type=str, help="Task name from task registry")
	parser.add_argument("--seed", default=DEFAULT_SEED, type=int, help="Deterministic random seed")
	parser.add_argument(
		"--all-tasks",
		action="store_true",
		help="Run reproducible baseline evaluation across easy, medium, and hard tasks",
	)
	parser.add_argument(
		"--max-steps",
		type=int,
		default=None,
		help="Optional cap on interaction steps (bounded by task max_steps)",
	)
	parser.add_argument(
		"--json-summary",
		action="store_true",
		help="When used with --all-tasks, print compact JSON summary",
	)
	return parser


def main() -> None:
	args = _build_arg_parser().parse_args()
	if args.all_tasks:
		summary = run_baseline_suite(seed=args.seed, max_steps=args.max_steps)
		if args.json_summary:
			print(json.dumps(summary, separators=(",", ":"), sort_keys=True))
			return

		for item in summary["tasks"]:
			print(
				f"[BASELINE] task={item['task']} "
				f"score={item['score']:.4f} "
				f"steps={item['steps']} "
				f"success={_bool_literal(item['success'])}"
			)

		print(
			f"[BASELINE_SUMMARY] seed={summary['seed']} "
			f"tasks={len(summary['tasks'])} "
			f"average_score={summary['average_score']:.4f}"
		)
		return

	run_episode(task=args.task, seed=args.seed, max_steps=args.max_steps, trace=True)


if __name__ == "__main__":
	main()