import os
import json
import tqdm
import argparse
from openai import OpenAI

# -----------------------------
#  CONFIGURATION
# -----------------------------
MODEL_NAME = "/home/mshahidul/readctrl_model/qwen3-32B_subclaims-extraction-8b_ctx"
API_URL = "http://localhost:8004/v1"
API_KEY = "EMPTY"

client = OpenAI(base_url=API_URL, api_key=API_KEY)

# -----------------------------
#  SUBCLAIM EXTRACTION PROMPT
# -----------------------------
def extraction_prompt(medical_text: str) -> str:
    return f"""
You are an expert medical annotator. Extract granular, factual subclaims.
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 output as:
[
  "subclaim 1",
  "subclaim 2",
  ...
]
"""

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

    final_prompt = extraction_prompt(medical_text)

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "user", "content": final_prompt}],
            max_tokens=1000,
            temperature=temperature,
            top_p=0.9,
        )
        res = response.choices[0].message.content.strip()
        
        # Handle cases where the model might include <think> tags or markdown code blocks
        if "</think>" in res:
            res = res.split("</think>")[-1].strip()
        
        if res.startswith("```json"):
            res = res.replace("```json", "").replace("```", "").strip()

        try:
            return json.loads(res)
        except:
            # Fallback if JSON parsing fails but some text is returned
            return [res]

    except Exception as e:
        print(f"API error for text snippet: {e}")
        return []


# ... (Configuration and extraction_prompt remain the same) ...

# -----------------------------
#  MAIN
# -----------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, 
                        default="/home/mshahidul/readctrl/data/synthetic_dataset_diff_labels/syn_data_with_gs_summary_en.json",
                        help="Path to input JSON file")
    parser.add_argument("--start_index", type=int, default=0,
                        help="Start index for processing")
    parser.add_argument("--end_index", type=int, default=-1,
                        help="End index for processing (exclusive). -1 = until end")

    args = parser.parse_args()

    SAVE_FOLDER = "/home/mshahidul/readctrl/data/extracting_subclaim"
    os.makedirs(SAVE_FOLDER, exist_ok=True)

    base_name = os.path.basename(args.input_file).replace(".json", "")
    OUTPUT_FILE = os.path.join(
        SAVE_FOLDER,
        f"subclaims_with_generated_{base_name}_{args.start_index}_{args.end_index}.json"
    )

    print(f"Loading {args.input_file}...")
    with open(args.input_file, "r") as f:
        data = json.load(f)

    total_items = len(data)
    start = args.start_index
    end = args.end_index if args.end_index != -1 else total_items
    work_items = data[start:end]

    result = []
    if os.path.exists(OUTPUT_FILE):
        try:
            with open(OUTPUT_FILE, "r") as f:
                result = json.load(f)
            print(f"Resuming. {len(result)} items already processed.")
        except:
            result = []

    # Using "index" or "id" as the unique identifier based on your JSON snippet
    existing_ids = {r.get("index") or r.get("id") for r in result}

    for item in tqdm.tqdm(work_items):
        # Handle different ID key names
        curr_id = item.get("index") if item.get("index") is not None else item.get("id")
        
        if curr_id in existing_ids:
            continue

        # 1. Process standard fields
        fulltext = item.get("fulltext", "")
        summary = item.get("summary", "")
        
        fulltext_sub = infer_subclaims(fulltext)
        summary_sub = infer_subclaims(summary)

        # 2. Process all generated texts (diff_label_texts)
        # We will create a mirror dictionary to store the subclaims
        diff_label_subclaims = {}
        generated_texts = item.get("diff_label_texts", {})

        for label, text in generated_texts.items():
            if text:
                diff_label_subclaims[label] = infer_subclaims(text)
            else:
                diff_label_subclaims[label] = []

        # 3. Build output object
        output_item = {
            "index": curr_id,
            "fulltext": fulltext,
            "fulltext_subclaims": fulltext_sub,
            "summary": summary,
            "summary_subclaims": summary_sub,
            "diff_label_texts": generated_texts,
            "diff_label_subclaims": diff_label_subclaims, # New field
            "readability_score": item.get("readability_score", None)
        }
        
        result.append(output_item)

        # Periodic save
        if len(result) % 10 == 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"Success! Results saved to: {OUTPUT_FILE}")