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 []

# -----------------------------
#  MAIN
# -----------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, 
                        default="/home/mshahidul/readctrl/data/classified_readability/classified_multiclinsum_test_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)

    # Output filename based on the source and range
    base_name = os.path.basename(args.input_file).replace(".json", "")
    OUTPUT_FILE = os.path.join(
        SAVE_FOLDER,
        f"subclaims_{base_name}_{args.start_index}_{args.end_index}.json"
    )

    # -----------------------------
    # Load data
    # -----------------------------
    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
    
    # Slice the data based on arguments
    work_items = data[start:end]

    print(f"Total records in file: {total_items}")
    print(f"Processing range: {start} → {end} (count={len(work_items)})")

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

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

    # -----------------------------
    # Process items
    # -----------------------------
    for item in tqdm.tqdm(work_items):
        _id = item.get("id")
        
        if _id in existing_ids:
            continue

        fulltext = item.get("fulltext", "")
        summary = item.get("summary", "")

        # Run inference for both fields
        fulltext_sub = infer_subclaims(fulltext)
        summary_sub = infer_subclaims(summary)

        # Build output object
        result.append({
            "id": _id,
            "fulltext": fulltext,
            "fulltext_subclaims": fulltext_sub,
            "summary": summary,
            "summary_subclaims": summary_sub,
            "readability_score": item.get("readability_score", None)
        })

        # Periodic save to prevent data loss
        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}")