import os
import re
import sys
import json
import time
import logging

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

START_TIME = time.time()
SOFT_LIMIT_SECONDS = 26 * 60
MODEL_ID = "."

import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    stream=sys.stdout,
    force=True,
)
log = logging.getLogger("iol-v1")


SYSTEM = (
    "You solve International Linguistics Olympiad problems by reasoning only from "
    "the supplied data. Read the instruction and examples carefully. Preserve every "
    "requested grammatical feature and output form. For translation, return the full "
    "translated form; for fill_blanks, return each missing form; for text_to_num, "
    "return digits; for num_to_text, return the complete written number form. Reason "
    "systematically and verify the inferred rule against the examples. Then write "
    "FINAL ANSWERS: exactly once, followed by one bare answer per line in query order. "
    "Do not number, quote, or explain the final answers. Always answer every item."
)


MATCH_SYSTEM = (
    "Solve this linguistic matching problem as a complete correspondence. Decompose "
    "recurring morphemes and compound words before matching. The numbered puzzle-language "
    "items must be answered in their original order using only option labels. When the "
    "number of labels equals the number of items, this is a one-to-one assignment: use "
    "every label exactly once. Resolve high-confidence pairs first, then use the global "
    "constraint for the remainder. Do not stop early; make a best guess for every item. "
    "Write FINAL ANSWERS: exactly once, followed by one uppercase label per line. Include "
    "no numbering, words, meanings, or commentary in the final section."
)


REPAIR_SYSTEM = (
    "Repair the final output of an International Linguistics Olympiad solution. Use the "
    "context, query, attempted solution, required count, and valid labels when supplied. "
    "Return exactly the required number of answers in original item order. For matching, "
    "return uppercase option labels only and use each label once when explicitly told the "
    "task is one-to-one. For text_to_num return digits; for num_to_text return complete "
    "written forms; for fill_blanks include every blank in order. Write FINAL ANSWERS: "
    "exactly once, then one bare answer per line. No reasoning or commentary."
)


MARKER_RE = re.compile(
    r"(?im)^\s*(?:[-*•]\s*)?(?:#{1,6}\s*)?(?:\*\*)?"
    r"final answers?\s*:?(?:\*\*)?\s*$"
)


def clean_answer_line(line):
    line = str(line).strip()
    if re.match(r"^#{1,6}\s+\S", line):
        return ""
    line = re.sub(r"^\s*\d+[.)]\s*", "", line)
    line = re.sub(r"^\s*[-*•]\s+", "", line).strip()
    if line.startswith("**") and line.endswith("**"):
        line = line[2:-2].strip()
    return line


def parse_answers(text):
    """Parse both a normal final block and repeated per-item FINAL ANSWERS blocks."""
    text = str(text)
    markers = list(MARKER_RE.finditer(text))
    if not markers:
        return []

    # Some models write one FINAL ANSWERS marker for every numbered item.
    if len(markers) > 1:
        answers = []
        for index, marker in enumerate(markers):
            end = markers[index + 1].start() if index + 1 < len(markers) else len(text)
            for line in text[marker.end():end].splitlines():
                candidate = clean_answer_line(line)
                if candidate:
                    answers.append(candidate)
                    break
        return answers

    answers = []
    for line in text[markers[0].end():].splitlines():
        if re.match(r"^\s*#{1,6}\s+\S", line):
            break
        candidate = clean_answer_line(line)
        if candidate:
            answers.append(candidate)
    return answers


def numbered_context_items(context):
    return list(dict.fromkeys(re.findall(r"(?m)^\s*(\d+)[.)]\s+\S", str(context))))


def option_labels(context):
    """Extract labels such as A.-R. from a matching context, preserving order."""
    labels = re.findall(r"(?m)^\s*([A-Z])\s*[.)]\s+\S", str(context))
    return list(dict.fromkeys(label.upper() for label in labels))


def count_query_items(query, context="", task_type=""):
    query = str(query).strip()
    task_type = str(task_type).strip()

    if task_type == "match_letters":
        items = numbered_context_items(context)
        if items:
            return len(items)

    ranges = re.findall(r"\((\d+)\s*[-–—]\s*(\d+)\)", query)
    if ranges:
        start, end = map(int, ranges[-1])
        if end >= start:
            return end - start + 1

    numbered = re.findall(r"(?m)^\s*\d+[.)]\s+\S", query)
    if numbered:
        return len(numbered)

    blanks = re.findall(r"\((\d+)\)", query)
    if blanks:
        return len(set(blanks))

    if ":" in query:
        body = query.split(":", 1)[1].strip()
        lines = [line.strip() for line in body.splitlines() if line.strip()]
        if len(lines) > 1:
            return len(lines)
        if len(lines) == 1:
            comma_items = [item.strip() for item in lines[0].split(",") if item.strip()]
            if len(comma_items) > 1:
                return len(comma_items)

    return 1


def normalize_matching_answers(answers):
    normalized = []
    for answer in answers:
        match = re.fullmatch(r"\s*([A-Za-z])(?:[.)])?\s*", str(answer))
        normalized.append(match.group(1).upper() if match else str(answer).strip().upper())
    return normalized


def matching_is_permutation(expected_n, labels):
    return len(labels) == expected_n and len(set(labels)) == expected_n


def valid_matching(answers, expected_n, labels):
    if len(answers) != expected_n:
        return False
    if not labels:
        return all(re.fullmatch(r"[A-Z]", answer or "") for answer in answers)
    if not all(answer in labels for answer in answers):
        return False
    if matching_is_permutation(expected_n, labels):
        return len(set(answers)) == expected_n and set(answers) == set(labels)
    return True


def complete_matching_permutation(answers, expected_n, labels):
    """Keep valid first occurrences and fill duplicates/missing slots with unused labels."""
    if not matching_is_permutation(expected_n, labels):
        return answers
    completed = []
    used = set()
    holes = []
    for answer in answers[:expected_n]:
        answer = str(answer).strip().upper()
        if answer in labels and answer not in used:
            completed.append(answer)
            used.add(answer)
        else:
            holes.append(len(completed))
            completed.append(None)
    while len(completed) < expected_n:
        holes.append(len(completed))
        completed.append(None)
    unused = [label for label in labels if label not in used]
    for position, label in zip(holes, unused):
        completed[position] = label
    return completed


def write_submission(rows):
    pd.DataFrame(rows, columns=["id", "pred"]).to_csv("submission.csv", index=False)


def generate_text(messages, max_new_tokens):
    input_ids = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)
    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(
        output[0][input_ids.shape[-1]:], skip_special_tokens=True
    ).strip()


def validate_submission(rows, source_df):
    errors = []
    if len(rows) != len(source_df):
        errors.append(f"row count {len(rows)} != {len(source_df)}")
    ids = [str(row["id"]) for row in rows]
    expected_ids = source_df["id"].astype(str).tolist()
    if ids != expected_ids:
        errors.append("IDs or row order differ from test.csv")
    if len(ids) != len(set(ids)):
        errors.append("duplicate IDs")
    for position, row in enumerate(rows):
        try:
            predictions = json.loads(row["pred"])
        except Exception as error:
            errors.append(f"row {position}: invalid JSON ({error})")
            continue
        source = source_df.iloc[position]
        expected_n = count_query_items(
            source["query"], source["context"], source.get("task_type", "")
        )
        if not isinstance(predictions, list):
            errors.append(f"row {position}: pred is not a list")
        elif len(predictions) != expected_n:
            errors.append(f"row {position}: {len(predictions)} != {expected_n} answers")
        elif not all(isinstance(answer, str) and answer != "" for answer in predictions):
            errors.append(f"row {position}: empty or non-string answer")
    return errors


log.info("Loading local model")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
    local_files_only=True,
).eval()
log.info("Model loaded in %.1fs", time.time() - START_TIME)

df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
log.info("Loaded %d test rows", len(df))

# Always keep a structurally complete CSV on disk.
submission_rows = []
for _, row in df.iterrows():
    expected_n = count_query_items(
        row["query"], row["context"], row.get("task_type", "")
    )
    labels = option_labels(row["context"]) if row.get("task_type", "") == "match_letters" else []
    fallback = labels[:expected_n] if len(labels) >= expected_n else ["?"] * expected_n
    submission_rows.append({
        "id": row["id"],
        "pred": json.dumps(fallback, ensure_ascii=False),
    })
write_submission(submission_rows)

for position, (_, row) in enumerate(df.iterrows()):
    task_type = row.get("task_type", "")
    expected_n = count_query_items(row["query"], row["context"], task_type)
    labels = option_labels(row["context"]) if task_type == "match_letters" else []
    one_to_one = task_type == "match_letters" and matching_is_permutation(expected_n, labels)
    elapsed = time.time() - START_TIME
    remaining = SOFT_LIMIT_SECONDS - elapsed
    if remaining > 10 * 60:
        main_budget = 1536
    elif remaining > 5 * 60:
        main_budget = 768
    else:
        main_budget = 384

    log.info(
        "row=%d/%d type=%s expected=%d labels=%d bijection=%s elapsed=%.1fs budget=%d",
        position + 1,
        len(df),
        task_type,
        expected_n,
        len(labels),
        one_to_one,
        elapsed,
        main_budget,
    )

    if task_type == "match_letters":
        user_content = (
            f"CONTEXT:\n{row['context'].strip()}\n\nQUERY:\n{row['query'].strip()}\n\n"
            f"Required answer count: {expected_n}.\n"
            f"Valid labels: {', '.join(labels) if labels else 'infer from context'}."
        )
        system_prompt = MATCH_SYSTEM
    else:
        user_content = (
            f"TASK TYPE: {task_type}\n\nCONTEXT:\n{row['context'].strip()}\n\n"
            f"QUERY:\n{row['query'].strip()}\n\nRequired answer count: {expected_n}."
        )
        system_prompt = SYSTEM

    try:
        raw_text = generate_text(
            [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content},
            ],
            max_new_tokens=main_budget,
        )
        log.info("row=%d primary_raw=%r", position + 1, raw_text)
        answers = parse_answers(raw_text)
        original_count = len(answers)
        if task_type == "match_letters":
            answers = normalize_matching_answers(answers)
            valid = valid_matching(answers, expected_n, labels)
        else:
            valid = len(answers) == expected_n and all(str(answer).strip() for answer in answers)

        log.info("row=%d primary parsed=%d valid=%s", position + 1, original_count, valid)

        if not valid and remaining > 3 * 60:
            repair_content = (
                f"TASK TYPE: {task_type}\nREQUIRED ANSWER COUNT: {expected_n}\n"
                f"ONE-TO-ONE MATCHING: {one_to_one}\n"
                f"VALID LABELS: {', '.join(labels)}\n\n"
                f"CONTEXT:\n{row['context'].strip()}\n\nQUERY:\n{row['query'].strip()}\n\n"
                f"ATTEMPTED SOLUTION:\n{raw_text}"
            )
            repair_text = generate_text(
                [
                    {"role": "system", "content": REPAIR_SYSTEM},
                    {"role": "user", "content": repair_content},
                ],
                max_new_tokens=768 if remaining > 8 * 60 else 384,
            )
            log.info("row=%d repair_raw=%r", position + 1, repair_text)
            repaired = parse_answers(repair_text)
            if task_type == "match_letters":
                repaired = normalize_matching_answers(repaired)
                repaired_valid = valid_matching(repaired, expected_n, labels)
            else:
                repaired_valid = len(repaired) == expected_n and all(
                    str(answer).strip() for answer in repaired
                )
            log.info(
                "row=%d repair parsed=%d valid=%s",
                position + 1,
                len(repaired),
                repaired_valid,
            )
            if repaired_valid or abs(len(repaired) - expected_n) < abs(len(answers) - expected_n):
                answers = repaired

        if task_type == "match_letters":
            answers = normalize_matching_answers(answers)
            if one_to_one and not valid_matching(answers, expected_n, labels):
                answers = complete_matching_permutation(answers, expected_n, labels)

        answers = [str(answer).strip() for answer in answers[:expected_n]]
        if len(answers) < expected_n:
            answers.extend(["?"] * (expected_n - len(answers)))
        answers = [answer if answer else "?" for answer in answers]

    except Exception as error:
        log.exception("row=%d failed: %s", position + 1, error)
        if task_type == "match_letters" and len(labels) >= expected_n:
            answers = labels[:expected_n]
        else:
            answers = ["?"] * expected_n
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    submission_rows[position] = {
        "id": row["id"],
        "pred": json.dumps(answers, ensure_ascii=False),
    }
    log.info(
        "row=%d id=%s final_answers=%s",
        position + 1,
        row["id"],
        json.dumps(answers, ensure_ascii=False),
    )
    write_submission(submission_rows)
    log.info("row=%d completed answers=%d/%d", position + 1, len(answers), expected_n)

errors = validate_submission(submission_rows, df)
if errors:
    for error in errors:
        log.error("validation: %s", error)
    raise RuntimeError("submission.csv validation failed")

log.info(
    "submission.csv valid rows=%d total_runtime=%.1fs",
    len(submission_rows),
    time.time() - START_TIME,
)