import os
# Set GPU environment variables
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import torch
from unsloth import FastLanguageModel
import json
import tqdm
import argparse



# -----------------------------
#  MODEL CACHE
# -----------------------------
_model_cache = {"model": None, "tokenizer": None}

def load_finetuned_model(model_path: str):
    """Load and cache your fine-tuned subclaim extraction model + tokenizer."""
    if _model_cache["model"] is not None:
        return _model_cache["model"], _model_cache["tokenizer"]

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_path,
        max_seq_length=8192,
        load_in_4bit=False,
        load_in_8bit=False,
        full_finetuning=False,
    )
    _model_cache["model"], _model_cache["tokenizer"] = model, tokenizer
    return model, tokenizer


# -----------------------------
#  SUBCLAIM EXTRACTION PROMPT
# -----------------------------
def extraction_prompt(medical_text: str) -> str:
    prompt = f"""
You are an expert medical annotator. Your task is to extract granular, factual subclaims from medical text.
A subclaim is the smallest standalone factual unit that can be independently verified.

Instructions:
1. Read the provided medical text.
2. Break it into clear, objective, atomic subclaims.
3. Each subclaim must come directly from the text.
4. Do not add, guess, or infer information.
5. Each subclaim should be short, specific, and verifiable.
6. Return ONLY a Python-style list of strings.

Medical Text:
{medical_text}

Return your output in JSON list format, like:
[
  "subclaim 1",
  "subclaim 2",
  ...
]
"""
    return prompt


# -----------------------------
#  INFERENCE FUNCTION
# -----------------------------
def infer_subclaims(medical_text: str,
                    model, 
                    tokenizer,
                    temperature: float = 0.2) -> list:
    
    if not medical_text or medical_text.strip() == "":
        return []

    prompt = extraction_prompt(medical_text)
    messages = [{"role": "user", "content": prompt}]

    chat_text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    inputs = tokenizer(chat_text, return_tensors="pt").to("cuda")

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=1024, # Increased to handle potentially longer list outputs
            temperature=temperature,
            top_p=0.9,
            top_k=10,
            do_sample=False,
        )

    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()

    # Remove thinking if model inserts `<think>`
    if "</think>" in output_text:
        output_text = output_text.split("</think>")[-1].strip()

    # Try to parse as JSON list, return raw text if parsing fails
    try:
        # Finding the start and end of the JSON list in case there is conversational filler
        start_idx = output_text.find('[')
        end_idx = output_text.rfind(']') + 1
        if start_idx != -1 and end_idx != -1:
            return json.loads(output_text[start_idx:end_idx])
        return output_text
    except Exception:
        return output_text


# -----------------------------
#  MAIN EXECUTION
# -----------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, required=True,
                        help="Path to the input JSON file containing medical texts.")
    args = parser.parse_args()
    
    INPUT_FILE = args.input_file
    file_name = os.path.basename(INPUT_FILE).split(".json")[0]
    SAVE_FOLDER = "/home/mshahidul/readctrl/data/extracting_subclaim"
    MODEL_PATH = "/home/mshahidul/readctrl_model/qwen3-32B_subclaims-extraction-8b_ctx"

    os.makedirs(SAVE_FOLDER, exist_ok=True)
    OUTPUT_FILE = os.path.join(SAVE_FOLDER, f"extracted_subclaims_{file_name}_en.json")

    # Load Model once
    model, tokenizer = load_finetuned_model(MODEL_PATH)

    # Load input dataset
    with open(INPUT_FILE, "r") as f:
        data = json.load(f)

    # Load existing results (resume mode)
    result = []
    if os.path.exists(OUTPUT_FILE):
        with open(OUTPUT_FILE, "r") as f:
            result = json.load(f)

    existing_ids = {item["id"] for item in result}

    # --------------------------------------------------------
    # PROCESS EACH MEDICAL TEXT (Fulltext AND Summary)
    # --------------------------------------------------------
    for item in tqdm.tqdm(data):
        if item.get("id") in existing_ids:
            continue

        # Extract from Fulltext
        fulltext_content = item.get("fulltext", "")
        fulltext_subclaims = infer_subclaims(fulltext_content, model, tokenizer)

        # Extract from Summary
        summary_content = item.get("summary", "")
        summary_subclaims = infer_subclaims(summary_content, model, tokenizer)

        result.append({
            "id": item.get("id"),
            "fulltext": fulltext_content,
            "fulltext_subclaims": fulltext_subclaims,
            "summary": summary_content,
            "summary_subclaims": summary_subclaims,
            "readability_score": item.get("readability_score", None)
        })

        # Save intermediate results
        if len(result) % 20 == 0:
            with open(OUTPUT_FILE, "w") as f:
                json.dump(result, f, indent=4, ensure_ascii=False)

    # Final save
    with open(OUTPUT_FILE, "w") as f:
        json.dump(result, f, indent=4, ensure_ascii=False)

    print(f"Extraction completed. Saved to {OUTPUT_FILE}")