#!/usr/bin/env python3
"""
Script to get OpenAlex Author IDs for EMBO members.

Uses three strategies based on available data:
1. ORCID (if available) - most reliable
2. Publications as context (if available)
3. Name + Affiliation + Subject Areas/Keywords (fallback)
"""

import os
import sys
import json
import asyncio
from pathlib import Path
from datetime import datetime
from typing import Optional, List

import pandas as pd
from dotenv import load_dotenv
from tqdm import tqdm

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from src.production_agent import disambiguate_author

load_dotenv()


def parse_publications(pub_str: str) -> List[str]:
    """
    Parse publications string into a list of publication titles.
    
    Args:
        pub_str: String containing publications (may be formatted in various ways)
    
    Returns:
        List of publication titles (up to 6)
    """
    if not pub_str or pd.isna(pub_str):
        return []
    
    pub_str = str(pub_str).strip()
    
    # Try to split by common delimiters
    # Publications might be separated by newlines, semicolons, or periods
    publications = []
    
    # Split by newline first
    if '\n' in pub_str:
        parts = pub_str.split('\n')
    # Then by semicolon
    elif ';' in pub_str:
        parts = pub_str.split(';')
    # Then by period (but be careful - periods are common in titles)
    else:
        # Try to split by ". " (period + space) which often indicates sentence end
        parts = pub_str.split('. ')
    
    for part in parts:
        part = part.strip()
        # Clean up common prefixes/suffixes
        part = part.strip('.,;')
        
        # Skip very short strings (likely not a publication)
        if len(part) > 20:
            publications.append(part)
        
        # Limit to 6 publications
        if len(publications) >= 6:
            break
    
    return publications[:6]


def determine_certainty_flag_and_reason(
    strategy: str,
    status: str,
    num_candidates: int,
    orcid_provided: Optional[str],
    orcid_matched: Optional[str],
    candidates: List[dict] = None
) -> tuple[str, str]:
    """
    Determine certainty flag and reason based on disambiguation results.
    
    Args:
        strategy: Strategy used (orcid, publications, name_affiliation_context)
        status: Result status (success, not_found, error, ambiguous)
        num_candidates: Number of candidates returned
        orcid_provided: ORCID provided in input
        orcid_matched: ORCID from matched author
        candidates: List of candidate dictionaries
    
    Returns:
        Tuple of (certainty_flag, reason)
    """
    if status == "error":
        return "error", "Error occurred during disambiguation"
    
    if status == "not_found":
        return "not_found", "No matching author found in OpenAlex"
    
    if status == "success":
        # Complete certainty: ORCID match with single candidate
        if strategy == "orcid" and num_candidates == 1:
            if orcid_provided and orcid_matched and orcid_provided == orcid_matched:
                return "complete_certainty", "ORCID exact match with single candidate"
            return "complete_certainty", "ORCID strategy with single candidate (very reliable)"
        
        # Check if we have candidates but status is success
        if num_candidates == 0:
            return "must_be_checked", "Success status but no candidates returned (schema mismatch or empty result)"
        
        # Medium certainty: Single candidate with publications or name+affiliation
        if num_candidates == 1:
            if strategy == "publications":
                return "medium_certainty", "Single candidate found using publications context"
            if strategy == "name_affiliation_context":
                return "medium_certainty", "Single candidate found using name, affiliation, and subject areas"
        
        # Multiple candidates - needs checking
        if num_candidates > 1:
            return "possible_multiple_ids", f"Multiple candidates returned ({num_candidates} candidates) - manual review needed"
        
        # Ambiguous status
        if status == "ambiguous":
            return "must_be_checked", "Ambiguous result - multiple plausible candidates"
    
    # Default fallback
    return "must_be_checked", f"Unknown reason - status: {status}, candidates: {num_candidates}"


def build_context_from_subjects_and_keywords(
    subject_areas: str,
    keywords: str
) -> Optional[str]:
    """
    Build context string from subject areas and keywords.
    
    Args:
        subject_areas: Subject areas string
        keywords: Keywords string
    
    Returns:
        Context string or None
    """
    context_parts = []
    
    if subject_areas and not pd.isna(subject_areas):
        context_parts.append(f"Research areas: {str(subject_areas)}")
    
    if keywords and not pd.isna(keywords):
        context_parts.append(f"Keywords: {str(keywords)}")
    
    if context_parts:
        return "\n".join(context_parts)
    
    return None


async def get_author_id_for_member(
    row: pd.Series,
    member_id: int
) -> dict:
    """
    Get OpenAlex author ID for a single EMBO member.
    
    Args:
        row: DataFrame row with member information
        member_id: Member identifier
    
    Returns:
        Dictionary with results
    """
    # Extract basic information
    first_name = row.get("mem_joi_PSUB__Person::PS_FirstNamePlusInitials", "")
    last_name = row.get("mem_joi_PSUB__Person::PS_LastName", "")
    affiliation = row.get("mem_joi_psub_add_INST__Institution#CurrentOnly::IN_FullName_Print__lct", "")
    
    # Determine strategy
    orcid = row.get("mem_joi_PSUB__Person::PS_ORCID", "")
    publications_str = row.get("mem_joi_joi_ELC__Elections::PUBLICATIONS", "")
    subject_areas = row.get("mem_joi_PSUB__Person::SC_SubjectAreas", "")
    keywords = row.get("mem_joi_PSUB__Person::SC_ML_OfficialKeywords", "")
    
    # Clean up values
    orcid = str(orcid).strip() if not pd.isna(orcid) else ""
    affiliation = str(affiliation).strip() if not pd.isna(affiliation) else None
    first_name = str(first_name).strip() if not pd.isna(first_name) else ""
    last_name = str(last_name).strip() if not pd.isna(last_name) else ""
    
    # Determine strategy and build context
    strategy = None
    context = None
    
    if orcid and orcid != "nan":
        strategy = "orcid"
    elif publications_str and not pd.isna(publications_str):
        strategy = "publications"
        publications = parse_publications(publications_str)
        if publications:
            context = "Known publications:\n" + "\n".join([f"- {p}" for p in publications])
    else:
        strategy = "name_affiliation_context"
        context = build_context_from_subjects_and_keywords(subject_areas, keywords)
    
    # Run disambiguation
    try:
        result = await disambiguate_author(
            first_name=first_name if first_name else None,
            last_name=last_name if last_name else None,
            affiliations=[affiliation] if affiliation else None,
            context=context,
            orcid=orcid if orcid and orcid != "nan" else None,
            find_email=False
        )
        
        # Extract author ID and all candidates from result
        author_id = None
        author_name = None
        author_orcid = None
        all_candidate_ids = []
        
        # Helper to normalize ID (extract from URL if needed)
        def normalize_id(id_str):
            if not id_str:
                return None
            id_str = str(id_str).strip()
            if "openalex.org/" in id_str:
                return id_str.split("openalex.org/")[-1].split("?")[0].split("&")[0]
            return id_str
        
        # Try new schema first
        candidates = result.get("author_candidates", [])
        
        if candidates:
            # Extract all candidate IDs
            for candidate in candidates:
                cand_id = candidate.get("author", {}).get("openalex_id", "")
                if not cand_id:
                    cand_id = candidate.get("openalex_author_id", "")
                normalized_id = normalize_id(cand_id)
                if normalized_id:
                    all_candidate_ids.append(normalized_id)
            
            # Get top candidate info
            if candidates:
                top_candidate = candidates[0]
                author_id = top_candidate.get("author", {}).get("openalex_id", "")
                if not author_id:
                    author_id = top_candidate.get("openalex_author_id", "")
                author_name = top_candidate.get("author", {}).get("name", "")
                if not author_name:
                    author_name = top_candidate.get("display_name", "")
                author_orcid = top_candidate.get("author", {}).get("orcid")
                if not author_orcid:
                    author_orcid = top_candidate.get("orcid")
        
        # Try old schema (if no candidates found)
        if not author_id and result.get("status") == "success":
            author_id = result.get("openalex_author_id", "")
            author_name = result.get("display_name", "")
            author_orcid = result.get("orcid")
            
            # If old schema has an ID, add it to candidates list
            if author_id:
                normalized_id = normalize_id(author_id)
                if normalized_id:
                    all_candidate_ids = [normalized_id]
                    author_id = normalized_id
        
        # Normalize top author ID
        if author_id:
            author_id = normalize_id(author_id)
        
        num_candidates = len(all_candidate_ids) if all_candidate_ids else len(candidates) if candidates else 0
        status = result.get("status", "")
        
        # Determine certainty flag and reason
        certainty_flag, check_reason = determine_certainty_flag_and_reason(
            strategy=strategy,
            status=status,
            num_candidates=num_candidates,
            orcid_provided=orcid if orcid and orcid != "nan" else None,
            orcid_matched=author_orcid,
            candidates=candidates
        )
        
        return {
            "member_id": member_id,
            "first_name": first_name,
            "last_name": last_name,
            "affiliation": affiliation,
            "strategy_used": strategy,
            "orcid_provided": orcid if orcid and orcid != "nan" else None,
            "publications_provided": len(parse_publications(publications_str)) if publications_str and not pd.isna(publications_str) else 0,
            "result": {
                "status": status,
                "openalex_author_id": author_id,
                "all_candidate_ids": all_candidate_ids if len(all_candidate_ids) > 1 else [],  # Only store if multiple
                "author_name": author_name,
                "author_orcid": author_orcid,
                "num_candidates": num_candidates,
                "certainty_flag": certainty_flag,
                "check_reason": check_reason
            },
            "metadata": result.get("_metadata", {})
        }
        
    except Exception as e:
        return {
            "member_id": member_id,
            "first_name": first_name,
            "last_name": last_name,
            "affiliation": affiliation,
            "strategy_used": strategy,
            "orcid_provided": orcid if orcid and orcid != "nan" else None,
            "publications_provided": len(parse_publications(publications_str)) if publications_str and not pd.isna(publications_str) else 0,
            "result": {
                "status": "error",
                "error": str(e),
                "openalex_author_id": None,
                "all_candidate_ids": [],
                "author_name": None,
                "author_orcid": None,
                "num_candidates": 0,
                "certainty_flag": "error",
                "check_reason": f"Error during processing: {str(e)}"
            },
            "metadata": {}
        }


async def process_embo_members(
    input_file: Path,
    output_file: Path = None,
    max_members: Optional[int] = None,
    start_from: int = 0,
    save_every: int = 1
):
    """
    Process EMBO members file and get OpenAlex author IDs.
    
    Args:
        input_file: Path to EMBO members Excel file
        output_file: Path to output Excel file (if None, auto-generates)
        max_members: Maximum number of members to process (None = all)
        start_from: Index to start from (for resuming)
        save_every: Save after every N members (default: 1 = save after each)
    """
    print(f"\n📂 Reading EMBO members file: {input_file}")
    df_full = pd.read_excel(input_file)
    
    print(f"✓ Loaded {len(df_full)} total members")
    
    # Determine output file
    if output_file is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = input_file.parent / f"emboplanet_allmem_with_openalex_ids_{timestamp}.xlsx"
    
    output_file = Path(output_file)
    
    # Check if output file already exists (for resuming)
    if output_file.exists():
        print(f"📂 Found existing output file: {output_file}")
        print(f"   Resuming from existing data...")
        df = pd.read_excel(output_file)
        
        # Ensure all required columns exist (add if missing)
        required_cols = {
            "openalex_author_id": "",
            "all_candidate_ids": "",
            "author_name": "",
            "author_orcid": "",
            "disambiguation_status": "",
            "strategy_used": "",
            "num_candidates": 0,
            "openalex_id_flag": "",
            "check_reason": "",
            "error_message": ""
        }
        for col, default_val in required_cols.items():
            if col not in df.columns:
                df[col] = default_val
        
        # Find where to resume (first row without openalex_author_id)
        if "openalex_author_id" in df.columns:
            completed = df["openalex_author_id"].notna() & (df["openalex_author_id"] != "")
            actual_start_from = completed.sum()
            if actual_start_from > start_from:
                start_from = actual_start_from
                print(f"   Found {start_from} already processed members")
                print(f"   Resuming from index {start_from}")
            else:
                print(f"   Starting from specified index {start_from}")
    else:
        # Create new dataframe with result columns
        df = df_full.copy()
        df["openalex_author_id"] = ""
        df["all_candidate_ids"] = ""  # Semicolon-separated list if multiple
        df["author_name"] = ""
        df["author_orcid"] = ""
        df["disambiguation_status"] = ""
        df["strategy_used"] = ""
        df["num_candidates"] = 0
        df["openalex_id_flag"] = ""
        df["check_reason"] = ""
        df["error_message"] = ""
    
    # Limit processing range if requested
    end_idx = len(df)
    if max_members:
        end_idx = min(start_from + max_members, len(df))
    
    df_to_process = df.iloc[start_from:end_idx]
    print(f"✓ Processing {len(df_to_process)} members (indices {start_from} to {end_idx-1})")
    
    # Process each member
    print(f"\n🔄 Processing members to get OpenAlex author IDs...\n")
    print(f"💾 Saving after every {save_every} member(s) to: {output_file}\n")
    
    results = []
    
    # Use iloc to get proper indexing
    for local_idx in range(len(df_to_process)):
        global_idx = start_from + local_idx
        row = df.iloc[global_idx]
        
        result = await get_author_id_for_member(row, global_idx)
        results.append(result)
        
        # Update dataframe
        df.at[global_idx, "openalex_author_id"] = result["result"].get("openalex_author_id", "") or ""
        
        # Store all candidate IDs if multiple
        all_candidate_ids = result["result"].get("all_candidate_ids", [])
        if len(all_candidate_ids) > 1:
            df.at[global_idx, "all_candidate_ids"] = "; ".join(all_candidate_ids)
        else:
            df.at[global_idx, "all_candidate_ids"] = ""
        
        df.at[global_idx, "author_name"] = result["result"].get("author_name", "") or ""
        df.at[global_idx, "author_orcid"] = result["result"].get("author_orcid", "") or ""
        df.at[global_idx, "disambiguation_status"] = result["result"].get("status", "")
        df.at[global_idx, "strategy_used"] = result.get("strategy_used", "")
        df.at[global_idx, "num_candidates"] = result["result"].get("num_candidates", 0)
        df.at[global_idx, "openalex_id_flag"] = result["result"].get("certainty_flag", "")
        df.at[global_idx, "check_reason"] = result["result"].get("check_reason", "")
        if result["result"].get("status") == "error":
            df.at[global_idx, "error_message"] = result["result"].get("error", "")
        
        # Save incrementally
        if (local_idx + 1) % save_every == 0 or local_idx == len(df_to_process) - 1:
            df.to_excel(output_file, index=False)
            if (local_idx + 1) % save_every == 0:
                print(f"   💾 Saved progress: {local_idx + 1}/{len(df_to_process)} members processed")
    
    # Save results
    if output_file is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = input_file.parent / f"emboplanet_allmem_with_openalex_ids_{timestamp}.xlsx"
    
    print(f"\n💾 Saving results to: {output_file}")
    df.to_excel(output_file, index=False)
    print(f"✓ Successfully saved {len(df)} members")
    
    # Print summary
    print(f"\n{'='*80}")
    print("SUMMARY")
    print(f"{'='*80}")
    
    total = len(results)
    success = sum(1 for r in results if r["result"]["status"] == "success" and r["result"]["openalex_author_id"])
    errors = sum(1 for r in results if r["result"]["status"] == "error")
    not_found = total - success - errors
    
    print(f"\nTotal processed: {total}")
    print(f"Successfully found: {success} ({success/total*100:.1f}%)")
    print(f"Not found: {not_found} ({not_found/total*100:.1f}%)")
    print(f"Errors: {errors} ({errors/total*100:.1f}%)")
    
    # Certainty flag breakdown
    print(f"\nCertainty Flag Breakdown:")
    certainty_counts = {}
    check_reason_counts = {}
    
    for r in results:
        flag = r["result"].get("certainty_flag", "unknown")
        certainty_counts[flag] = certainty_counts.get(flag, 0) + 1
        
        # Count check reasons for "must_be_checked" cases
        if flag == "must_be_checked":
            reason = r["result"].get("check_reason", "Unknown reason")
            check_reason_counts[reason] = check_reason_counts.get(reason, 0) + 1
    
    flag_descriptions = {
        "complete_certainty": "Complete certainty (ORCID match, single candidate)",
        "medium_certainty": "Medium certainty (single candidate, publications/name context)",
        "must_be_checked": "Must be checked (ambiguous or low confidence)",
        "possible_multiple_ids": "Possible multiple IDs (multiple candidates returned)",
        "not_found": "Not found",
        "error": "Error during processing"
    }
    
    for flag, count in sorted(certainty_counts.items(), key=lambda x: x[1], reverse=True):
        desc = flag_descriptions.get(flag, flag)
        print(f"  - {flag}: {count} ({count/total*100:.1f}%) - {desc}")
    
    # Show check reasons breakdown
    if check_reason_counts:
        print(f"\nTop reasons for 'must_be_checked' flag:")
        for reason, count in sorted(check_reason_counts.items(), key=lambda x: x[1], reverse=True)[:5]:
            print(f"  - {reason}: {count} cases")
    
    # Strategy breakdown
    print(f"\nStrategy breakdown:")
    strategy_counts = {}
    for r in results:
        strategy = r.get("strategy_used", "unknown")
        strategy_counts[strategy] = strategy_counts.get(strategy, 0) + 1
    
    for strategy, count in strategy_counts.items():
        success_count = sum(1 for r in results if r.get("strategy_used") == strategy and r["result"]["status"] == "success" and r["result"]["openalex_author_id"])
        print(f"  - {strategy}: {count} members ({success_count} successful, {success_count/count*100:.1f}%)")
    
    # Save detailed results JSON
    json_output = output_file.with_suffix('.json')
    print(f"\n💾 Saving detailed results to: {json_output}")
    with open(json_output, 'w') as f:
        json.dump({
            "timestamp": datetime.now().isoformat(),
            "total_members": total,
            "summary": {
                "success": success,
                "not_found": not_found,
                "errors": errors
            },
            "certainty_breakdown": certainty_counts,
            "strategy_breakdown": strategy_counts,
            "results": results
        }, f, indent=2)
    
    print(f"\n{'='*80}")
    print("✅ COMPLETE")
    print(f"{'='*80}\n")


def main():
    """CLI entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Get OpenAlex Author IDs for EMBO members"
    )
    parser.add_argument(
        "--input",
        type=str,
        default="data/emboplanet_allmem_2026-01-07.xlsx",
        help="Input Excel file path"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output Excel file path (auto-generated if not provided)"
    )
    parser.add_argument(
        "-n", "--max-members",
        type=int,
        default=None,
        help="Maximum number of members to process (for testing)"
    )
    parser.add_argument(
        "--start-from",
        type=int,
        default=0,
        help="Index to start from (for resuming)"
    )
    parser.add_argument(
        "--save-every",
        type=int,
        default=1,
        help="Save after every N members (default: 1 = save after each member)"
    )
    
    args = parser.parse_args()
    
    input_file = Path(args.input)
    if not input_file.exists():
        print(f"❌ Error: Input file not found: {input_file}")
        sys.exit(1)
    
    output_file = Path(args.output) if args.output else None
    
    asyncio.run(process_embo_members(
        input_file=input_file,
        output_file=output_file,
        max_members=args.max_members,
        start_from=args.start_from,
        save_every=args.save_every
    ))


if __name__ == "__main__":
    main()
