import os
import re
import sys

_SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
_SUBMIT_ROOT = os.path.dirname(os.path.dirname(_SCRIPTS_DIR))
if _SCRIPTS_DIR not in sys.path:
    sys.path.insert(0, _SCRIPTS_DIR)
if _SUBMIT_ROOT not in sys.path:
    sys.path.insert(0, _SUBMIT_ROOT)

_DIRECTION_PATTERN = re.compile(
    r"situated to the (left of|right of|in front of|behind|below|above) the",
    re.IGNORECASE,
)

_SKIP_DIRECTIONS = {"below", "above"}

_DIRECTION_CANONICAL = {
    "left of":     "left",
    "right of":    "right",
    "in front of": "front",
    "behind":      "behind",
}

_DIRECTION_TO_PHRASE = {
    "left":        "left of",
    "right":       "right of",
    "in_front_of": "in front of",
    "behind":      "behind",
}

_FACING_PROMPT_TEMPLATE = (
    "In the image, there is a {object_name}. "
    "Which direction is the {object_name} facing from the camera's perspective? "
    "Choose exactly one and output only that choice: "
    "left, right, toward the camera, away from the camera."
)

_OBJECT_EXTRACTION_SYSTEM = (
    "From the following sentence, identify the target object that is being acted upon, "
    "placed, or referred to as the main subject of interest. "
    "Output only the object name, nothing else."
)


def should_remap(question):
    match = _DIRECTION_PATTERN.search(question)
    if match is None:
        return False, None, None
    phrase = match.group(1).lower()
    if phrase in _SKIP_DIRECTIONS:
        return False, None, None
    canonical = _DIRECTION_CANONICAL.get(phrase)
    return True, canonical, phrase


def _parse_facing_direction(raw_output):
    text = raw_output.lower()
    if "away from" in text or "backward" in text:
        return "facing_away_from_camera"
    if "toward the camera" in text or "towards the camera" in text:
        return "facing_toward_camera"
    if "forward" in text:
        return "facing_toward_camera"
    has_left = "left" in text
    has_right = "right" in text
    if has_left and not has_right:
        return "facing_left"
    if has_right and not has_left:
        return "facing_right"
    if "toward" in text or "towards" in text:
        return "facing_toward_camera"
    return None


def _map_direction(orig_canonical, facing_key, direction_map):
    if orig_canonical is None or facing_key is None:
        return None
    facing_entry = direction_map.get(facing_key, {})
    return facing_entry.get(orig_canonical)


def _replace_direction_in_question(question, orig_phrase, new_canonical):
    new_phrase = _DIRECTION_TO_PHRASE.get(new_canonical)
    if new_phrase is None:
        return question
    old_pattern = f"to the {orig_phrase} the"
    new_pattern = f"to the {new_phrase} the"
    return question.replace(old_pattern, new_pattern)


def _extract_object_from_question(question, clf_kwargs):
    m = re.search(r"there is a (.+?)\.", question, re.IGNORECASE)
    if m:
        return m.group(1).strip()

    import torch
    from lm_classifier import _apply_chat_template

    first_sentence = (question.split(".")[0] + ".") if "." in question else question

    clf_model = clf_kwargs["model"]
    clf_tokenizer = clf_kwargs["tokenizer"]
    first_device = next(clf_model.parameters()).device

    messages = [
        {"role": "system", "content": _OBJECT_EXTRACTION_SYSTEM},
        {"role": "user", "content": f"Sentences: {first_sentence}"},
    ]
    text = _apply_chat_template(clf_tokenizer, messages, enable_thinking=False)
    inputs = clf_tokenizer([text], return_tensors="pt").to(first_device)

    with torch.no_grad():
        generated_ids = clf_model.generate(**inputs, max_new_tokens=16, do_sample=False)
    trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, generated_ids)]
    raw = clf_tokenizer.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0].strip()

    if "</think>" in raw:
        raw = raw.split("</think>", 1)[1].strip()
    if not raw or raw == "[]":
        raw = "object"
    return raw


def run_context_with_remap(question, image_path, depth_path, model_kwargs, clf_kwargs, direction_map):
    import torch
    from robobrain_runner import run_robobrain
    from evaluation import _extract_first_point

    do_remap, orig_dir_canonical, orig_dir_phrase = should_remap(question)
    if not do_remap:
        return run_robobrain(question, image_path, depth_path, model_kwargs, LM_classify="context")

    answer1 = run_robobrain(question, image_path, depth_path, model_kwargs, LM_classify="context")
    coord1_tuple, _ = _extract_first_point(answer1)
    coord1_str = f"({coord1_tuple[0]}, {coord1_tuple[1]})" if coord1_tuple else ""

    object_name = _extract_object_from_question(question, clf_kwargs)
    torch.cuda.empty_cache()

    obj_label = object_name.strip() or "object"
    dir_prompt = _FACING_PROMPT_TEMPLATE.format(object_name=obj_label)
    dir_answer = run_robobrain(dir_prompt, image_path, depth_path, model_kwargs, add_think_override=False)
    facing_key = _parse_facing_direction(dir_answer)

    new_dir = _map_direction(orig_dir_canonical, facing_key, direction_map)
    if new_dir is None:
        return coord1_str or answer1

    new_question = _replace_direction_in_question(question, orig_dir_phrase, new_dir)

    answer2 = run_robobrain(new_question, image_path, depth_path, model_kwargs, LM_classify="context")
    coord2_tuple, _ = _extract_first_point(answer2)
    coord2_str = f"({coord2_tuple[0]}, {coord2_tuple[1]})" if coord2_tuple else ""

    parts = [p for p in [coord1_str, coord2_str] if p]
    return " ".join(parts) if parts else (answer1 or answer2)