#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import csv
import re
import json
import time
import datetime
import urllib.request
import urllib.parse
import openpyxl
from bs4 import BeautifulSoup
from concurrent.futures import ThreadPoolExecutor, as_completed

script_dir = os.path.dirname(os.path.abspath(__file__))
base_dir = os.path.dirname(script_dir)

OLLAMA_URL = "http://192.168.221.20:11434/api/generate"
MODEL_NAME = "medgemma:27b"
LOG_FILE_PATH = os.path.join(script_dir, "ollama_calls.log")


def log_ollama_call(path_code, attempt, duration_sec, status, site="", ttype="", tissue="", diagnosis="", error=None):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    file_exists = os.path.exists(LOG_FILE_PATH) and os.path.getsize(LOG_FILE_PATH) > 0
    
    try:
        with open(LOG_FILE_PATH, "a", newline="", encoding="utf-8-sig") as f:
            writer = csv.writer(f)
            if not file_exists:
                writer.writerow(["timestamp", "path_code", "attempt", "duration_sec", "status", "site", "type", "tissue", "diagnosis", "error"])
            
            writer.writerow([
                timestamp,
                path_code,
                attempt,
                round(duration_sec, 3),
                status,
                site,
                ttype,
                tissue,
                diagnosis,
                str(error) if error else ""
            ])
    except Exception as e:
        pass


def extract_plain_text(html_content):
    if not html_content or not isinstance(html_content, str):
        return ""
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # Remove script and style tags
    for s in soup(['script', 'style', 'head', 'meta', 'link']):
        s.decompose()
        
    text = soup.get_text(separator=' ').strip()
    return re.sub(r'\s+', ' ', text)[:3000]


def enforce_site_rules(site, tissue, diagnosis):
    tissue_lower = str(tissue).lower()

    # Rule: If tissue is Nasopharynx or Nasal cavity, site is NOT cranial (should be other).
    # However, sinonasal IS cranial.
    if 'sinonasal' in tissue_lower:
        site = 'cranial'
    elif 'nasopharynx' in tissue_lower or 'nasal cavity' in tissue_lower:
        if site == 'cranial':
            site = 'other'

    return site


def extract_with_medgemma_pure(html_content, path_code='', retries=5):
    text_content = extract_plain_text(html_content)

    prompt = f"""You are an expert medical pathology analysis assistant.
Analyze the following pathology report text for specimen {path_code} and extract key clinical information.

Report Text:
{text_content}

Tasks:
1. "site": Anatomical site classification. Must be EXACTLY one of: "cranial", "spinal", or "other".
   - "cranial": Brain, cerebrum, cerebellum, brainstem, skull, cranial dura, pituitary, CP angle, orbit, sinonasal tissue, head, scalp, etc.
   - "NOTE on Nasal cavity / Nasopharynx": If tissue origin is "Nasopharynx" or "Nasal cavity" alone, it is NOT cranial (classify as "other"). However, "sinonasal" IS cranial.
   - "spinal": Spine, spinal cord, vertebra, lumbar/thoracic/cervical/sacral level, epidural, disc, ligamentum flavum, etc.
   - "other": All other body organs or unspecified sites.

2. "type": Pathological nature. Must be EXACTLY one of: "tumor" or "other".
   - "tumor": Neoplasms, carcinomas, adenomas, gliomas, meningiomas, schwannomas, metastases, lipomas, polyps, cysts, or neoplastic findings.
   - "other": Non-neoplastic findings (e.g. inflammation, hematoma, abscess, necrosis, normal tissue, thrombus).

3. "tissue": Tissue origin / specimen site string in English (e.g., "Brain, cerebrum", "Soft tissue, L3", "Spine, L4-L5 level", "Bone, T12 vertebra", "Spinal cord, T7-T9").
   - CRITICAL: Always keep and preserve the specific spine level if specified in the report (e.g. "L3", "L4/5", "T7-9", "C4-5", "T12", "S1-S2", "C1-C2"). Do not strip or omit spine levels!

4. "diagnosis": Concise histological / pathological diagnosis statement in English.

Return ONLY a valid JSON object matching this schema:
{{
  "site": "cranial|spinal|other",
  "type": "tumor|other",
  "tissue": "tissue origin (including exact spine level if specified)",
  "diagnosis": "pathological diagnosis"
}}
"""

    payload = {
        "model": MODEL_NAME,
        "prompt": prompt,
        "stream": False,
        "format": "json"
    }

    for attempt in range(1, retries + 1):
        start_time = time.time()
        try:
            req = urllib.request.Request(
                OLLAMA_URL,
                data=json.dumps(payload).encode('utf-8'),
                headers={'Content-Type': 'application/json'}
            )
            with urllib.request.urlopen(req, timeout=90) as resp:
                elapsed = time.time() - start_time
                res_data = json.loads(resp.read().decode('utf-8'))
                resp_text = res_data.get('response', '')
                parsed = json.loads(resp_text)
                
                site = str(parsed.get('site', '')).strip().lower()
                ttype = str(parsed.get('type', '')).strip().lower()
                tissue = str(parsed.get('tissue', '')).strip()
                diagnosis = str(parsed.get('diagnosis', '')).strip()

                if site not in ['cranial', 'spinal', 'other']:
                    site = 'other'
                if ttype not in ['tumor', 'other']:
                    ttype = 'other'

                site = enforce_site_rules(site, tissue, diagnosis)

                log_ollama_call(
                    path_code=path_code,
                    attempt=attempt,
                    duration_sec=elapsed,
                    status="SUCCESS",
                    site=site,
                    ttype=ttype,
                    tissue=tissue,
                    diagnosis=diagnosis
                )

                return site, ttype, tissue, diagnosis

        except Exception as e:
            elapsed = time.time() - start_time
            print(f"Warning: MedGemma API attempt {attempt}/{retries} error for {path_code}: {e}")
            log_ollama_call(
                path_code=path_code,
                attempt=attempt,
                duration_sec=elapsed,
                status="ERROR",
                error=e
            )
            if attempt < retries:
                time.sleep(3 * attempt)

    return "other", "other", "", ""


def process_single_row(args):
    idx, chart_no, path_code, html_content, total_count = args
    print(f"[{idx+1}/{total_count}] Extracting PathCode {path_code} with MedGemma...")
    site, ttype, tissue, diagnosis = extract_with_medgemma_pure(html_content, path_code)
    print(f"  -> [{idx+1}/{total_count}] {path_code}: site={site}, type={ttype}, tissue={tissue!r}, diag={diagnosis!r}")
    return idx, [chart_no, path_code, site, ttype, tissue, diagnosis]


def process_extraction_medgemma(input_excel=None, output_csv=None, limit=0, prefix_filter=None, max_workers=4):
    if input_excel is None:
        input_excel = os.path.join(script_dir, 'registry_pathologyreport.xlsx')
        if not os.path.exists(input_excel) and os.path.exists('registry_pathologyreport.xlsx'):
            input_excel = 'registry_pathologyreport.xlsx'

    if output_csv is None:
        if prefix_filter:
            output_csv = os.path.join(script_dir, f'registry_pathologyreport_medgemma_{prefix_filter}.csv')
        else:
            output_csv = os.path.join(script_dir, 'registry_pathologyreport_medgemma_all.csv')

    print(f"Opening Excel file: {input_excel} ...")
    wb = openpyxl.load_workbook(input_excel, read_only=True)
    ws = wb.active

    headers = None
    chart_idx = None
    path_idx = None
    html_idx = None

    raw_rows = []
    print(f"Reading Excel rows (prefix filter: {prefix_filter!r}, limit: {limit})...")
    for i, row in enumerate(ws.iter_rows()):
        vals = [cell.value for cell in row]
        if headers is None:
            headers = vals
            chart_idx = headers.index('ChartNo')
            path_idx = headers.index('PathCode')
            html_idx = headers.index('html')
            continue

        path_code = str(vals[path_idx]).strip() if vals[path_idx] else ""
        if prefix_filter:
            if not path_code.startswith(prefix_filter):
                continue

        raw_rows.append(vals)
        if limit and len(raw_rows) >= limit:
            break

    total_count = len(raw_rows)
    print(f"Total rows to extract: {total_count} (max_workers: {max_workers})...")

    tasks_args = [
        (idx, vals[chart_idx], str(vals[path_idx]).strip() if vals[path_idx] else "", vals[html_idx] or "", total_count)
        for idx, vals in enumerate(raw_rows)
    ]

    results_dict = {}
    completed_count = 0

    # Write CSV header initially
    with open(output_csv, 'w', newline='', encoding='utf-8-sig') as f:
        writer = csv.writer(f)
        writer.writerow(['病歷號', '病理號', 'site', 'type', '組織由來', 'pathological diagnosis'])

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_idx = {executor.submit(process_single_row, arg): arg[0] for arg in tasks_args}
        
        for future in as_completed(future_to_idx):
            idx, res_row = future.result()
            results_dict[idx] = res_row
            completed_count += 1

            # Periodically dump progress every 50 rows
            if completed_count % 50 == 0 or completed_count == total_count:
                print(f"Progress: {completed_count}/{total_count} rows completed ({completed_count/total_count*100:.2f}%)...")

    print(f"Writing all {total_count} rows in original order to CSV: {output_csv} ...")
    with open(output_csv, 'w', newline='', encoding='utf-8-sig') as f:
        writer = csv.writer(f)
        writer.writerow(['病歷號', '病理號', 'site', 'type', '組織由來', 'pathological diagnosis'])
        for idx in range(total_count):
            writer.writerow(results_dict[idx])

    print(f"Successfully exported {total_count} rows to {output_csv}")


if __name__ == '__main__':
    lim = int(sys.argv[1]) if len(sys.argv) > 1 and sys.argv[1].isdigit() else 0
    pref = sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] != 'None' else None
    out = sys.argv[3] if len(sys.argv) > 3 else None
    workers = int(sys.argv[4]) if len(sys.argv) > 4 and sys.argv[4].isdigit() else 4
    process_extraction_medgemma(limit=lim, prefix_filter=pref, output_csv=out, max_workers=workers)
