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

Strategy:
1. For members < 2008: Use disambiguation agent with name + context
2. For members >= 2008 with publications: Find work IDs, extract author IDs from works
3. For members >= 2008 without publications: Use disambiguation agent

Outputs:
- Excel file with work_ids and author_ids as comma-separated lists
- JSON file with detailed responses
"""

import os
import sys
import json
import re
import asyncio
from pathlib import Path
from typing import Optional, List, Dict, Tuple

import pandas as pd
from tqdm import tqdm
import pyalex
from pyalex import Works, Authors

# Add src directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
from production_agent import disambiguate_author  # noqa: E402

# Initialize pyalex
email = os.getenv('OPENALEX_API_KEY')
if email:
    pyalex.config.email = email


def parse_publications(publications_text: str) -> List[Dict[str, str]]:
    """
    Parse publications text into structured list.
    
    Attempts to extract:
    - Title
    - DOI
    - Journal
    - Year
    - Authors (for matching)
    
    Args:
        publications_text: Raw publications text from Excel
        
    Returns:
        List of publication dictionaries
    """
    if not publications_text or pd.isna(publications_text):
        return []
    
    publications = []
    lines = publications_text.split('\n')
    
    for line in lines:
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        
        # Skip numbered list markers
        line = re.sub(r'^\d+\.\s*', '', line)
        
        pub = {
            'raw_text': line,
            'title': None,
            'doi': None,
            'journal': None,
            'year': None,
            'authors': None
        }
        
        # Extract DOI
        doi_match = re.search(r'doi[:\s]*([10]\.[0-9]+/[^\s\)]+)', line, re.IGNORECASE)
        if doi_match:
            pub['doi'] = doi_match.group(1)
        
        # Extract year (4-digit year, typically near the end)
        year_match = re.search(r'\b(19|20)\d{2}\b', line)
        if year_match:
            pub['year'] = int(year_match.group(0))
        
        # Try to extract title (usually between authors and journal/year)
        # This is heuristic - titles are often in parentheses or before journal names
        title_patterns = [
            r'\((\d{4})\)\s+([^\.]+?)(?:\.|,|$)',
            r'([A-Z][^\(\)]+?)(?:\s*\(|\s*\.|$)',
        ]
        for pattern in title_patterns:
            match = re.search(pattern, line)
            if match:
                potential_title = match.group(1) if len(match.groups()) > 0 else match.group(0)
                if len(potential_title) > 10:  # Reasonable title length
                    pub['title'] = potential_title.strip()
                    break
        
        # Extract journal (often after title, before year)
        journal_patterns = [
            r'\.\s+([A-Z][^,\.]+?)(?:,|\s+\d{4})',
            r'([A-Z][A-Za-z\s&]+?)(?:\s+\d{4})',
        ]
        for pattern in journal_patterns:
            match = re.search(pattern, line)
            if match:
                pub['journal'] = match.group(1).strip()
                break
        
        publications.append(pub)
    
    return publications


def search_work_by_doi(doi: str) -> Optional[Dict]:
    """Search for a work by DOI in OpenAlex."""
    try:
        # Normalize DOI
        if not doi.startswith('10.'):
            return None
        
        # Search by DOI
        works = Works().filter(doi=doi).get(per_page=1)
        if works:
            return works[0]
    except Exception as e:
        print(f"  ⚠️  Error searching DOI {doi}: {e}")
    return None


def search_work_by_title(title: str, year: Optional[int] = None) -> Optional[Dict]:
    """Search for a work by title in OpenAlex."""
    try:
        # Clean title for search
        search_title = title[:200]  # Limit length
        
        # Search by title
        filter_params = {"title.search": search_title}
        if year:
            filter_params["publication_year"] = year
        
        works = Works().filter(**filter_params).get(per_page=5)
        
        # Try to find best match
        if works:
            # Simple heuristic: prefer exact title match or highest cited
            for work in works:
                work_title = work.get('title', '').lower()
                if work_title == title.lower():
                    return work
            # Return most cited if no exact match
            return max(works, key=lambda w: w.get('cited_by_count', 0))
    except Exception as e:
        print(f"  ⚠️  Error searching title '{title[:50]}...': {e}")
    return None


def extract_author_id_from_work(work: Dict, first_name: str, last_name: str) -> Optional[str]:
    """
    Extract OpenAlex author ID from work's author list by matching name.
    
    Args:
        work: OpenAlex work object
        first_name: First name with initials (e.g., "J. M.")
        last_name: Last name
        
    Returns:
        OpenAlex author ID (without URL prefix) or None
    """
    if not work or 'authorships' not in work:
        return None
    
    # Normalize names for matching
    last_name_lower = last_name.lower().strip()
    first_name_parts = [p.strip().lower() for p in first_name.split() if p.strip()]
    
    # Try to match author
    for authorship in work.get('authorships', []):
        author = authorship.get('author', {})
        if not author:
            continue
        
        author_name = author.get('display_name', '')
        if not author_name:
            continue
        
        # Split author name
        name_parts = author_name.split()
        if len(name_parts) < 2:
            continue
        
        author_last = name_parts[-1].lower()
        author_first_parts = [p.lower() for p in name_parts[:-1]]
        
        # Check last name match
        if author_last != last_name_lower:
            continue
        
        # Check first name/initials match
        # Match if first letter of first parts match
        if first_name_parts:
            first_letters = [p[0] for p in first_name_parts if p]
            author_first_letters = [p[0] for p in author_first_parts if p]
            
            if first_letters and author_first_letters:
                # Check if initials match
                if all(fl in author_first_letters or any(fl == af[0] for af in author_first_parts) 
                       for fl in first_letters):
                    # Extract author ID
                    author_id = author.get('id', '')
                    if author_id:
                        # Return just the ID part (without URL)
                        return author_id.split('/')[-1]
        
        # Also try exact match on last name if first name is very short/ambiguous
        if len(first_name_parts) == 0 or (len(first_name_parts) == 1 and len(first_name_parts[0]) == 1):
            author_id = author.get('id', '')
            if author_id:
                return author_id.split('/')[-1]
    
    return None


async def find_author_ids_from_publications(
    publications: List[Dict],
    first_name: str,
    last_name: str
) -> Tuple[List[str], List[str], Dict]:
    """
    Find OpenAlex work IDs and author IDs from publications.
    
    Args:
        publications: List of parsed publication dictionaries
        first_name: First name with initials
        last_name: Last name
        
    Returns:
        Tuple of (work_ids, author_ids, details_dict)
    """
    work_ids = []
    author_ids = []
    details = {
        'works_found': 0,
        'authors_found': 0,
        'errors': []
    }
    
    for pub in publications:
        work = None
        
        # Try DOI first (most reliable)
        if pub.get('doi'):
            work = search_work_by_doi(pub['doi'])
            if work:
                details['works_found'] += 1
        
        # Fallback to title search
        if not work and pub.get('title'):
            work = search_work_by_title(pub['title'], pub.get('year'))
            if work:
                details['works_found'] += 1
        
        if work:
            work_id = work.get('id', '').split('/')[-1]
            work_ids.append(work_id)
            
            # Extract author ID from work
            author_id = extract_author_id_from_work(work, first_name, last_name)
            if author_id:
                if author_id not in author_ids:
                    author_ids.append(author_id)
                details['authors_found'] += 1
        else:
            details['errors'].append(f"Work not found: {pub.get('title', pub.get('raw_text', 'Unknown'))[:50]}")
    
    return work_ids, author_ids, details


def search_author_by_orcid(orcid: str) -> Optional[Dict]:
    """Search for author by ORCID in OpenAlex."""
    try:
        # Normalize ORCID
        orcid_clean = str(orcid).strip()
        if not orcid_clean:
            return None
        
        # Remove URL prefix if present
        orcid_clean = orcid_clean.replace("https://orcid.org/", "").replace("http://orcid.org/", "")
        
        # Search by ORCID
        authors = Authors().filter(orcid=orcid_clean).get(per_page=1)
        if authors:
            author = authors[0]
            return {
                'openalex_id': author.get('id', '').split('/')[-1],
                'display_name': author.get('display_name', ''),
                'orcid': author.get('orcid', ''),
                'works_count': author.get('works_count', 0)
            }
    except Exception as e:
        print(f"  ⚠️  Error searching ORCID {orcid}: {e}")
    return None


async def process_member(
    row: pd.Series,
    year_col: str,
    first_name_col: str,
    last_name_col: str,
    publications_col: str,
    orcid_col: str,
    affiliation_col: str,
    country_col: str,
    keywords_col: str,
    subject_areas_col: str
) -> Dict:
    """
    Process a single member to find their OpenAlex ID.
    
    Strategy:
    1. If ORCID exists: Use agent with ORCID + verify with publications if available
    2. If publications available: Use agent with publications as context + extract work IDs
    3. Otherwise: Use agent with name + affiliations + keywords
    
    Returns:
        Dictionary with results including:
        - openalex_author_ids: List of author IDs (comma-separated string)
        - openalex_work_ids: List of work IDs (comma-separated string)
        - strategy: Strategy used
        - status: success/error
        - details: Detailed information
    """
    result = {
        'openalex_author_ids': '',
        'openalex_work_ids': '',
        'strategy': '',
        'status': 'pending',
        'certainty_flag': 'pending',
        'check_reason': '',
        'details': {}
    }
    
    # Get member data
    first_name = str(row.get(first_name_col, '')).strip()
    last_name = str(row.get(last_name_col, '')).strip()
    publications_text = row.get(publications_col)
    orcid = row.get(orcid_col)
    affiliation = row.get(affiliation_col)
    country = row.get(country_col)
    keywords = row.get(keywords_col)
    subject_areas = row.get(subject_areas_col)
    
    # Build context for agent
    context_parts = []
    
    # Add publications as context if available
    if publications_text and not pd.isna(publications_text):
        context_parts.append(f"Publications:\n{publications_text}")
    
    # Add keywords and subject areas
    if keywords and not pd.isna(keywords):
        context_parts.append(f"Keywords: {keywords}")
    if subject_areas and not pd.isna(subject_areas):
        context_parts.append(f"Subject Areas: {subject_areas}")
    
    context = '\n\n'.join(context_parts) if context_parts else None
    
    # Build affiliations list
    affiliations_list = []
    if affiliation and not pd.isna(affiliation):
        affiliations_list.append(str(affiliation))
    if country and not pd.isna(country):
        affiliations_list.append(str(country))
    
    # Determine strategy based on available data
    has_orcid = orcid and not pd.isna(orcid)
    has_publications = publications_text and not pd.isna(publications_text)
    
    if has_orcid and has_publications:
        result['strategy'] = 'orcid_with_publications_verification'
    elif has_orcid:
        result['strategy'] = 'orcid_with_agent_verification'
    elif has_publications:
        result['strategy'] = 'agent_with_publications'
    else:
        result['strategy'] = 'agent_with_context'
    
    # Call disambiguation agent
    try:
        agent_result = await disambiguate_author(
            first_name=first_name,
            last_name=last_name,
            affiliations=affiliations_list if affiliations_list else None,
            context=context,
            orcid=str(orcid) if has_orcid else None,
            find_email=False,
            max_iterations=15
        )
        
        # Extract author IDs from agent result
        author_ids_from_agent = []
        if agent_result.get('status') == 'success':
            # Handle both old and new schema
            if 'author_candidates' in agent_result:
                for candidate in agent_result['author_candidates']:
                    author_id = candidate.get('openalex_author_id', '')
                    if author_id:
                        # Extract ID from URL if needed
                        if '/' in author_id:
                            author_id = author_id.split('/')[-1]
                        if author_id not in author_ids_from_agent:
                            author_ids_from_agent.append(author_id)
            elif 'openalex_author_id' in agent_result:
                author_id = agent_result['openalex_author_id']
                if author_id:
                    if '/' in author_id:
                        author_id = author_id.split('/')[-1]
                    author_ids_from_agent.append(author_id)
        
        result['status'] = agent_result.get('status', 'error')
        result['details']['agent_result'] = agent_result
        
        # If we have publications, also try to extract work IDs for verification
        work_ids = []
        author_ids_from_works = []
        if has_publications:
            publications = parse_publications(str(publications_text))
            if publications:
                work_ids_found, author_ids_found, work_details = await find_author_ids_from_publications(
                    publications, first_name, last_name
                )
                work_ids.extend(work_ids_found)
                author_ids_from_works.extend(author_ids_found)
                result['details']['work_extraction'] = work_details
        
        # Combine author IDs from agent and works
        all_author_ids = []
        all_author_ids.extend(author_ids_from_agent)
        for aid in author_ids_from_works:
            if aid not in all_author_ids:
                all_author_ids.append(aid)
        
        result['openalex_author_ids'] = ','.join(all_author_ids) if all_author_ids else ''
        result['openalex_work_ids'] = ','.join(work_ids) if work_ids else ''
        
        # Determine certainty flag
        if not all_author_ids:
            result['certainty_flag'] = 'not_found'
            result['check_reason'] = 'No author IDs found by agent or work extraction'
        elif has_orcid:
            # ORCID was provided
            if len(all_author_ids) == 1:
                if author_ids_from_works and author_ids_from_works[0] == all_author_ids[0]:
                    result['certainty_flag'] = 'complete_certainty'
                    result['check_reason'] = f'ORCID confirmed by {len(work_ids)} works'
                else:
                    result['certainty_flag'] = 'medium_certainty'
                    result['check_reason'] = 'ORCID found but not fully verified by works'
            else:
                result['certainty_flag'] = 'must_be_checked'
                result['check_reason'] = f'ORCID provided but multiple IDs found: {all_author_ids}'
        elif has_publications:
            # Publications were used
            unique_ids = list(set(all_author_ids))
            if len(unique_ids) == 1:
                if len(work_ids) >= 3 and author_ids_from_works:
                    result['certainty_flag'] = 'complete_certainty'
                    result['check_reason'] = f'Single consistent author ID across {len(work_ids)} works'
                else:
                    result['certainty_flag'] = 'medium_certainty'
                    result['check_reason'] = f'Single ID from agent with {len(work_ids)} works'
            else:
                result['certainty_flag'] = 'possible_multiple_ids'
                result['check_reason'] = f'Multiple author IDs found: {unique_ids}'
        else:
            # Only name + context used (old members or no publications)
            if len(all_author_ids) == 1:
                result['certainty_flag'] = 'medium_certainty'
                result['check_reason'] = 'Single ID from agent (no publications/ORCID for verification)'
            else:
                result['certainty_flag'] = 'possible_multiple_ids'
                result['check_reason'] = f'Multiple IDs from agent: {all_author_ids}'
                
    except Exception as e:
        result['status'] = 'error'
        result['certainty_flag'] = 'error'
        result['check_reason'] = f'Exception: {str(e)}'
        result['details'] = {'error': str(e)}
    
    return result


async def process_all_members(
    input_file: Path,
    output_file: Optional[Path] = None,
    json_output_file: Optional[Path] = None,
    max_members: Optional[int] = None,
    start_from: int = 0,
    save_every: int = 10
):
    """
    Process all EMBO members to find their OpenAlex IDs.
    
    Args:
        input_file: Path to cleaned Excel file
        output_file: Path to output Excel file
        json_output_file: Path to JSON output file
        max_members: Maximum number of members to process (None = all)
        start_from: Index to start from (for resuming)
        save_every: Save progress every N members
    """
    print(f"\n{'='*80}")
    print("PROCESSING EMBO MEMBERS FOR OPENALEX IDs")
    print(f"{'='*80}\n")
    
    # Load data
    print(f"📂 Loading data from: {input_file}")
    df = pd.read_excel(input_file)
    print(f"✓ Loaded {len(df)} members\n")
    
    # Column names
    year_col = "mem_joi_PSUB__Person::PS_EMBO_Year_Membership"
    first_name_col = "mem_joi_PSUB__Person::PS_FirstNamePlusInitials"
    last_name_col = "mem_joi_PSUB__Person::PS_LastName"
    publications_col = "mem_joi_joi_ELC__Elections::PUBLICATIONS"
    orcid_col = "mem_joi_PSUB__Person::PS_ORCID"
    affiliation_col = "mem_joi_psub_add_INST__Institution#CurrentOnly::IN_FullName_Print__lct"
    country_col = "mem_joi_psub_add_COU__Country#CurrentOnly::Country_Name"
    keywords_col = "mem_joi_PSUB__Person::SC_ML_OfficialKeywords"
    subject_areas_col = "mem_joi_PSUB__Person::SC_SubjectAreas"
    
    # Initialize output columns if they don't exist
    output_cols = [
        'openalex_author_ids',
        'openalex_work_ids',
        'openalex_strategy',
        'openalex_status',
        'openalex_certainty_flag',
        'openalex_check_reason'
    ]
    
    for col in output_cols:
        if col not in df.columns:
            df[col] = ''
    
    # Determine range
    end_idx = min(start_from + max_members, len(df)) if max_members else len(df)
    members_to_process = df.iloc[start_from:end_idx]
    
    print(f"Processing members {start_from} to {end_idx-1} ({len(members_to_process)} members)\n")
    
    # Process each member
    all_results = []
    
    for idx, (row_idx, row) in enumerate(tqdm(members_to_process.iterrows(), total=len(members_to_process), desc="Processing members")):
        try:
            result = await process_member(
                row,
                year_col,
                first_name_col,
                last_name_col,
                publications_col,
                orcid_col,
                affiliation_col,
                country_col,
                keywords_col,
                subject_areas_col
            )
            
            # Update DataFrame
            df.at[row_idx, 'openalex_author_ids'] = result['openalex_author_ids']
            df.at[row_idx, 'openalex_work_ids'] = result['openalex_work_ids']
            df.at[row_idx, 'openalex_strategy'] = result['strategy']
            df.at[row_idx, 'openalex_status'] = result['status']
            df.at[row_idx, 'openalex_certainty_flag'] = result['certainty_flag']
            df.at[row_idx, 'openalex_check_reason'] = result['check_reason']
            
            # Store detailed result
            result['row_index'] = int(row_idx)
            result['first_name'] = str(row.get(first_name_col, ''))
            result['last_name'] = str(row.get(last_name_col, ''))
            all_results.append(result)
            
            # Save incrementally
            if (idx + 1) % save_every == 0:
                if output_file:
                    df.to_excel(output_file, index=False)
                if json_output_file:
                    with open(json_output_file, 'w') as f:
                        json.dump(all_results, f, indent=2, default=str)
                print(f"\n💾 Progress saved at member {idx + 1}")
        
        except Exception as e:
            print(f"\n❌ Error processing member {row_idx}: {e}")
            df.at[row_idx, 'openalex_status'] = 'error'
            df.at[row_idx, 'openalex_certainty_flag'] = 'error'
            df.at[row_idx, 'openalex_check_reason'] = f'Exception: {str(e)}'
    
    # Final save
    if output_file:
        print(f"\n💾 Saving final results to: {output_file}")
        df.to_excel(output_file, index=False)
    
    if json_output_file:
        print(f"💾 Saving detailed JSON to: {json_output_file}")
        with open(json_output_file, 'w') as f:
            json.dump(all_results, f, indent=2, default=str)
    
    # Summary
    print(f"\n{'='*80}")
    print("PROCESSING COMPLETE")
    print(f"{'='*80}\n")
    
    processed_df = df.iloc[start_from:end_idx]
    print(f"Processed: {len(processed_df)} members")
    print("\nStatus distribution:")
    print(processed_df['openalex_status'].value_counts())
    print("\nCertainty flag distribution:")
    print(processed_df['openalex_certainty_flag'].value_counts())
    print("\nStrategy distribution:")
    print(processed_df['openalex_strategy'].value_counts())


def main():
    """CLI entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Find OpenAlex IDs for all EMBO members"
    )
    parser.add_argument(
        "--input",
        type=str,
        default="data/emboplanet_allmem_cleaned_2026-01-07.xlsx",
        help="Input cleaned Excel file"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="data/emboplanet_allmem_with_openalex_ids.xlsx",
        help="Output Excel file"
    )
    parser.add_argument(
        "--json-output",
        type=str,
        default="output/embo_members_openalex_results.json",
        help="Output JSON file with detailed results"
    )
    parser.add_argument(
        "--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=10,
        help="Save progress every N members"
    )
    
    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)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    
    json_output_file = Path(args.json_output)
    json_output_file.parent.mkdir(parents=True, exist_ok=True)
    
    asyncio.run(process_all_members(
        input_file=input_file,
        output_file=output_file,
        json_output_file=json_output_file,
        max_members=args.max_members,
        start_from=args.start_from,
        save_every=args.save_every
    ))


if __name__ == "__main__":
    main()
