"""
Fetch dosage recommendation data for medicines listed in a CSV file.

Usage:
    python scripts/fetch_dosage_data.py [--csv PATH] [--batch-size N] [--api-url URL] [--age AGE] [--gender GENDER]

Defaults:
    --csv       scripts/query_result_2026-04-07t11_22_2023564794405_30.csv
    --batch-size 5
    --api-url   http://localhost:8001/api/v1/drugs/dosage/recommendation
    --age       (optional)
    --gender    (optional)
"""

import argparse
import csv
import json
import time
from pathlib import Path

import requests


def fetch_dosage_batch(
    medications: list[str],
    api_url: str,
    age: int | None = None,
    gender: str | None = None,
) -> dict:
    """Call the recommendation API for a batch of medicines. Returns {medicine_name: data_dict}."""
    payload = {"medications": medications}
    if age is not None:
        payload["age"] = age
    if gender is not None:
        payload["gender"] = gender

    resp = requests.post(api_url, json=payload, params={"ai_enhanced": True}, timeout=300)
    resp.raise_for_status()
    body = resp.json()

    results = {}
    for item in body.get("recommendationData", {}).get("data", []):
        name = item.get("medicine_name", "")
        results[name] = item
    return results


def fetch_all_dosages(
    medicine_names: list[str],
    api_url: str,
    age: int | None = None,
    gender: str | None = None,
    batch_size: int = 5,
    delay: float = 0.5,
) -> dict:
    """Fetch dosage data for all medicines in controlled batches. Returns {name: data_dict}."""
    all_results = {}
    total = len(medicine_names)

    for i in range(0, total, batch_size):
        batch = medicine_names[i : i + batch_size]
        batch_num = i // batch_size + 1
        total_batches = (total + batch_size - 1) // batch_size
        print(f"  Batch {batch_num}/{total_batches}: {len(batch)} medicines ...")

        try:
            results = fetch_dosage_batch(batch, api_url, age, gender)
            all_results.update(results)
            print(f"    Got {len(results)} results")
        except Exception as e:
            print(f"    ERROR: {e}")
            for name in batch:
                if name not in all_results:
                    all_results[name] = None

        if i + batch_size < total:
            time.sleep(delay)

    return all_results


def process_csv(
    csv_path: str,
    api_url: str,
    age: int | None = None,
    gender: str | None = None,
    batch_size: int = 5,
    drug_name_column: str = "drug_name",
    dosage_column: str = "Dosage",
):
    """Read medicines from CSV, fetch dosage data, and write results back to the same CSV."""
    path = Path(csv_path)
    if not path.exists():
        raise FileNotFoundError(f"CSV not found: {path}")

    # Read CSV
    with open(path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        fieldnames = list(reader.fieldnames)
        rows = list(reader)

    if drug_name_column not in fieldnames:
        raise ValueError(f"Column '{drug_name_column}' not found. Available: {fieldnames}")

    # Collect medicine names
    medicine_names = [row[drug_name_column].strip() for row in rows if row[drug_name_column].strip()]
    print(f"Found {len(medicine_names)} medicines in '{csv_path}'")

    # Fetch dosage data
    all_results = fetch_all_dosages(medicine_names, api_url, age, gender, batch_size)

    # Add Dosage column if not present
    if dosage_column not in fieldnames:
        fieldnames.append(dosage_column)

    # Update rows
    found = 0
    for row in rows:
        name = row[drug_name_column].strip()
        data = all_results.get(name)
        if data:
            row[dosage_column] = json.dumps(data, ensure_ascii=False)
            found += 1
        else:
            row[dosage_column] = ""

    # Write back to same CSV
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"\nDone! Updated '{csv_path}'")
    print(f"  Found: {found}/{len(medicine_names)}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fetch dosage data for medicines in a CSV")
    parser.add_argument(
        "--csv",
        default="scripts/query_result_2026-04-07t11_22_2023564794405_30.csv",
        help="Path to CSV file",
    )
    parser.add_argument("--batch-size", type=int, default=5, help="Medicines per API call")
    parser.add_argument("--api-url", default="http://localhost:8001/api/v1/drugs/dosage/recommendation")
    parser.add_argument("--age", type=int, default=None, help="Patient age (optional)")
    parser.add_argument("--gender", default=None, help="Patient gender (optional)")
    args = parser.parse_args()

    process_csv(args.csv, args.api_url, args.age, args.gender, args.batch_size)