"""
Works Search Agent

Finds OpenAlex works (papers) based on title and author information.
Uses a hybrid approach: direct PyAlex search with strict validation,
falling back to an AI agent for uncertain cases.
"""

import asyncio
import os
from typing import Optional, Dict, Any
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Import PyAlex
from pyalex import Works
import pyalex

# Configure PyAlex
pyalex.config.email = os.getenv("OPENALEX_API_KEY", "")
pyalex.config.max_retries = 3
pyalex.config.retry_backoff_factor = 0.1
pyalex.config.retry_http_codes = [429, 500, 503]


def normalize_string(s: str) -> str:
    """
    Normalize a string for comparison by:
    - Converting to lowercase
    - Removing extra whitespace
    - Removing common punctuation
    """
    if not s:
        return ""
    
    # Convert to lowercase
    s = s.lower()
    
    # Remove common punctuation
    for char in '.,;:!?()[]{}"\'-–—':
        s = s.replace(char, ' ')
    
    # Collapse multiple spaces into one
    s = ' '.join(s.split())
    
    return s.strip()


def titles_match(title1: str, title2: str, threshold: float = 0.9) -> bool:
    """
    Check if two titles match after normalization.
    
    Args:
        title1: First title
        title2: Second title
        threshold: Similarity threshold (0-1)
    
    Returns:
        True if titles match closely enough
    """
    norm1 = normalize_string(title1)
    norm2 = normalize_string(title2)
    
    # Exact match after normalization
    if norm1 == norm2:
        return True
    
    # Check if one is a substring of the other (handles subtitle variations)
    if len(norm1) > 20 and len(norm2) > 20:  # Only for longer titles
        if norm1 in norm2 or norm2 in norm1:
            return True
    
    # Could add fuzzy matching here if needed (e.g., using difflib)
    # For now, require exact match after normalization
    return False


def author_in_work(
    author_last_name: str, work: Dict[str, Any]
) -> tuple[bool, Optional[str], Optional[str]]:
    """
    Check if an author with the given last name is in the work's author list.
    
    Args:
        author_last_name: Author's last name to search for
        work: OpenAlex work object
    
    Returns:
        Tuple of (found: bool, openalex_id: str, orcid: str)
    """
    if not author_last_name or not work:
        return False, None, None
    
    norm_last_name = normalize_string(author_last_name)
    authorships = work.get('authorships', [])
    
    for authorship in authorships:
        author = authorship.get('author', {})
        display_name = author.get('display_name', '')
        
        # Normalize and check if last name is in display name
        norm_display_name = normalize_string(display_name)
        
        if norm_last_name in norm_display_name:
            # Found! Extract OpenAlex ID and ORCID
            author_id = author.get('id', '')
            orcid = author.get('orcid', '')
            return True, author_id, orcid
    
    return False, None, None


def search_work_direct(
    title: str,
    author_last_name: Optional[str] = None,
    year: Optional[int] = None
) -> Optional[Dict[str, Any]]:
    """
    Search for a work using direct PyAlex queries with strict validation.
    
    Args:
        title: Paper title
        author_last_name: Author's last name (optional but recommended)
        year: Publication year (optional)
    
    Returns:
        Work data dict if found and validated, None otherwise
    """
    if not title:
        return None
    
    try:
        # Remove commas from title (PyAlex interprets them as separators)
        clean_title = title.replace(',', '')
        
        # Search OpenAlex by title
        works = Works()
        search = works.search_filter(title=clean_title)
        
        # Add year filter if provided
        if year:
            search = search.filter(publication_year=year)
        
        results = search.get()
        
        if not results:
            return None
        
        # Check each result for strict validation
        for work in results[:5]:  # Check top 5 results
            work_title = work.get('title', '')
            
            # Validate: Title must match
            if not titles_match(title, work_title):
                continue
            
            # If author provided, validate author is in work
            author_found = True
            author_id = None
            author_orcid = None
            
            if author_last_name:
                author_found, author_id, author_orcid = author_in_work(
                    author_last_name, work
                )
                
                # If author not found, skip this result
                if not author_found:
                    continue
            
            # Found a match! Return work data
            return {
                "openalex_id": work.get('id', ''),
                "doi": work.get('doi', ''),
                "title": work_title,
                "author_openalex_id": author_id,
                "author_orcid": author_orcid,
                "author_found_in_work": author_found,
                "match_confidence": "direct"
            }
        
        # No validated match found
        return None
        
    except Exception as e:
        print(f"Error in direct search: {e}")
        return None


async def search_work_with_agent(
    title: str,
    author_last_name: Optional[str] = None,
    author_first_name: Optional[str] = None,
    year: Optional[int] = None,
    max_iterations: int = 10
) -> Optional[Dict[str, Any]]:
    """
    Search for a work using the AI agent when direct search fails.

    This uses Claude Agent SDK with OpenAlex MCP tools and web_search
    to intelligently find the work, handling fuzzy matches and alternate
    titles.

    Args:
        title: Paper title
        author_last_name: Author's last name
        author_first_name: Author's first name
        year: Publication year
        max_iterations: Maximum agent iterations

    Returns:
        Work data dict if found, None otherwise
    """
    # Import here to avoid circular dependencies
    import sys
    import json
    from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions
    from pathlib import Path

    # Add src to path for imports
    src_path = Path(__file__).parent
    if str(src_path) not in sys.path:
        sys.path.insert(0, str(src_path))

    # Import MCP options and schema
    from openalex_mcp import OPENALEX_MCP_OPTIONS
    from schemas.works_search_result import WORKS_SEARCH_SCHEMA

    # Build search query for the agent
    query_parts = [f"Find the OpenAlex work for the paper titled: '{title}'"]

    if author_last_name:
        author_name = (f"{author_first_name} {author_last_name}"
                       if author_first_name else author_last_name)
        query_parts.append(f"by author: {author_name}")

    if year:
        query_parts.append(f"published in: {year}")

    query = " ".join(query_parts)

    # System prompt for the agent
    system_prompt = f"""You are a works search expert using OpenAlex MCP tools.

Your task: Find the OpenAlex work ID for a given paper.

**CRITICAL VALIDATION RULES:**
1. Title must match after normalization (case-insensitive, ignore punctuation)
2. If author specified, they MUST appear in the work's author list
3. Set author_found_in_work flag accurately

**Available Tools:**
- search_works_by_title: Search papers by title
- search_works_by_doi: Get paper by DOI
- search_works_by_title_and_author: Combined title + author search
- get_work_details: Get complete work information
- WebSearch: Find DOI if needed

**Process:**
1. Use search_works_by_title or search_works_by_title_and_author
2. Validate title matches
3. Find author '{author_last_name if author_last_name else "N/A"}'
   in authorships
4. Extract author's OpenAlex ID and ORCID
5. Return structured JSON

**RETURN ONLY** valid JSON matching this schema:
{{
  "status": "success" | "not_found",
  "work": {{
    "openalex_id": "W...",
    "doi": "10.../...",
    "title": "...",
    "author_openalex_id": "A..." or null,
    "author_orcid": "0000-..." or null,
    "author_found_in_work": true/false,
    "match_confidence": "agent"
  }}
}}"""

    try:
        # Configure agent options (following production_agent.py pattern)
        agent_options = ClaudeAgentOptions(
            mcp_servers=OPENALEX_MCP_OPTIONS.mcp_servers,
            allowed_tools=OPENALEX_MCP_OPTIONS.allowed_tools + ["Skill"],
            system_prompt=system_prompt,
            cwd=os.path.dirname(os.path.abspath(__file__)),  # src/ directory (where .claude/skills lives)
            setting_sources=["project"],  # Load skills from .claude/skills
            output_format=WORKS_SEARCH_SCHEMA,  # Enforce structured output
        )

        final_json = None

        # Use ClaudeSDKClient with async context manager
        async with ClaudeSDKClient(options=agent_options) as client:
            # Send the query
            await client.query(query)

            # Receive and process response
            all_text = ""
            async for message in client.receive_response():
                # Collect text for JSON extraction
                if hasattr(message, 'content'):
                    for block in message.content:
                        block_type = type(block).__name__
                        if (block_type == 'TextBlock'
                                and hasattr(block, 'text')):
                            all_text += block.text
                        elif (hasattr(block, 'type')
                              and block.type == "text"
                              and hasattr(block, 'text')):
                            all_text += block.text

            # Extract JSON from the collected text
            text = all_text.strip()

            # Remove markdown code blocks if present
            if text.startswith('```json'):
                text = text[7:]
            elif text.startswith('```'):
                text = text[3:]
            if text.endswith('```'):
                text = text[:-3]
            text = text.strip()

            # Try to parse as JSON
            try:
                final_json = json.loads(text)
            except json.JSONDecodeError:
                # Try to extract JSON with regex
                import re
                json_match = re.search(r'\{.*\}', text, re.DOTALL)
                if json_match:
                    try:
                        final_json = json.loads(json_match.group())
                    except Exception:
                        pass

            # Return result if valid
            if final_json and final_json.get('status') == 'success':
                if final_json.get('work'):
                    final_json['work']['match_confidence'] = 'agent'
                return final_json

            return None

    except Exception as e:
        print(f"Error in agent search: {e}")
        import traceback
        traceback.print_exc()
        return None


async def search_work(
    title: str,
    author_last_name: Optional[str] = None,
    author_first_name: Optional[str] = None,
    year: Optional[int] = None,
    max_iterations: int = 10
) -> Dict[str, Any]:
    """
    Search for an OpenAlex work using a hybrid approach.
    
    First attempts direct PyAlex search with strict validation.
    Falls back to AI agent if no confident match is found.
    
    Args:
        title: Paper title (required)
        author_last_name: Author's last name (recommended)
        author_first_name: Author's first name (optional)
        year: Publication year (optional)
        max_iterations: Maximum agent iterations (default: 10)
    
    Returns:
        Dict with status and work data:
        {
            "status": "success" | "not_found" | "error",
            "work": {
                "openalex_id": "W1234567890",
                "doi": "10.1234/example",
                "title": "Full work title",
                "author_openalex_id": "A5023888391",
                "author_orcid": "0000-0001-2345-6789",
                "author_found_in_work": true,
                "match_confidence": "direct" | "agent"
            }
        }
    """
    try:
        # Step 1: Try direct PyAlex search with strict validation
        result = search_work_direct(title, author_last_name, year)
        
        if result:
            return {
                "status": "success",
                "work": result
            }
        
        # Step 2: Fall back to agent search
        result = await search_work_with_agent(
            title, author_last_name, author_first_name, year, max_iterations
        )
        
        if result and result.get("status") == "success":
            return result
        
        # Not found
        return {
            "status": "not_found",
            "work": None,
            "message": f"No validated match found for title: {title[:50]}..."
        }
        
    except Exception as e:
        return {
            "status": "error",
            "work": None,
            "error": str(e)
        }


def main():
    """CLI entry point for works search."""
    import sys
    import json
    
    if len(sys.argv) < 2:
        print("Usage: search-work --title \"Paper Title\" "
              "[--author-last \"LastName\"] [--year 2023]")
        sys.exit(1)
    
    # Simple argument parsing
    title = None
    author_last = None
    year = None
    
    i = 1
    while i < len(sys.argv):
        if sys.argv[i] == '--title' and i + 1 < len(sys.argv):
            title = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--author-last' and i + 1 < len(sys.argv):
            author_last = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--year' and i + 1 < len(sys.argv):
            year = int(sys.argv[i + 1])
            i += 2
        else:
            i += 1
    
    if not title:
        print("Error: --title is required")
        sys.exit(1)
    
    # Run search
    result = asyncio.run(search_work(title, author_last, year=year))
    
    # Print result
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()
