#!/usr/bin/env python3
"""Translate psychiatry_cases diagnosis names to Korean using the Anthropic API."""

import sqlite3
import json
import time
import sys
import urllib.request

DB_PATH = "/home/kim/homeclaw/.smallclaw/databases/psychiatry_cases.db"
API_URL = "http://localhost:11434/v1/chat/completions"
MODEL = "kimi-k2.6:cloud"


def translate_batch(items):
    prompt = """Translate the following DSM-5 psychiatric disorder names from English to Korean.
Rules:
- Use official/standard Korean DSM-5 terminology where established
- Keep acronyms like ADHD, PTSD, OCD as-is if commonly used in Korean
- Keep the parenthetical clarifications when present
- Each line has format: NUMBER|ENGLISH_NAME
- Return ONLY the translations in format: NUMBER|KOREAN_NAME
- No explanations, no notes

"""
    for id_, name in items:
        prompt += f"{id_}|{name}\n"

    data = json.dumps({
        "model": MODEL,
        "messages": [
            {"role": "system", "content": "You are a professional Korean-English translator specializing in psychiatry and DSM-5 terminology. Use accurate, standard Korean psychiatric terms."},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.2,
    }).encode("utf-8")

    req = urllib.request.Request(
        API_URL,
        data=data,
        headers={"Content-Type": "application/json"},
    )

    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            result = json.loads(resp.read().decode("utf-8"))
            return result["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"API error: {e}", file=sys.stderr, flush=True)
        return None


def main():
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()

    c.execute("SELECT id, diagnosis_en FROM cases WHERE (diagnosis_ko = '' OR diagnosis_ko IS NULL) AND diagnosis_en != '' ORDER BY id")
    rows = c.fetchall()

    if not rows:
        print("All diagnoses translated!")
        conn.close()
        return

    print(f"Translating {len(rows)} diagnosis names...", flush=True)

    batch_size = 30
    total_translated = 0

    for i in range(0, len(rows), batch_size):
        batch = rows[i:i + batch_size]
        print(f"\nBatch {i // batch_size + 1}/{(len(rows) + batch_size - 1) // batch_size} ({len(batch)} items)...", flush=True)

        result = translate_batch(batch)
        if not result:
            print("Translation failed, skipping batch", flush=True)
            continue

        updates = []
        for line in result.strip().split("\n"):
            line = line.strip()
            if "|" not in line:
                continue
            parts = line.split("|", 1)
            try:
                id_ = int(parts[0].strip())
                ko = parts[1].strip().strip("*").strip()
                updates.append((ko, id_))
                total_translated += 1
            except (ValueError, IndexError):
                print(f"  Skipping: {line}", flush=True)

        if updates:
            c.executemany("UPDATE cases SET diagnosis_ko = ? WHERE id = ?", updates)
            conn.commit()

        print(f"  Translated so far: {total_translated}", flush=True)
        time.sleep(2)

    c.execute("SELECT COUNT(*) FROM cases WHERE diagnosis_ko != '' AND diagnosis_ko IS NOT NULL")
    done = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM cases WHERE diagnosis_ko = '' OR diagnosis_ko IS NULL")
    remaining = c.fetchone()[0]

    print(f"\nDone! Translated: {done}, Remaining: {remaining}", flush=True)

    print("\n--- Sample translations ---")
    c.execute("SELECT diagnosis_en, diagnosis_ko FROM cases WHERE diagnosis_ko != '' AND diagnosis_ko IS NOT NULL ORDER BY RANDOM() LIMIT 10")
    for row in c.fetchall():
        print(f"  EN: {row[0]}")
        print(f"  KO: {row[1]}")
        print()

    conn.close()


if __name__ == "__main__":
    main()