#!/usr/bin/env python3
"""
Production Author Disambiguation Agent

Uses Claude Agent SDK with embedded MCP pattern:
- OpenAlex MCP tools for author search and verification
- Web search for additional context
- JSON-only structured output
- Iteration limits

Architecture:
- Uses ClaudeSDKClient from claude_agent_sdk
- MCP server is embedded (not via stdio)
- Clean separation of concerns via MCP module
- Async execution throughout
"""

import os
import sys
import json
import asyncio
from dotenv import load_dotenv

from claude_agent_sdk import (
    ClaudeSDKClient,
    ClaudeAgentOptions,
)

# Add src directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from prompts.direct_system_prompt import DIRECT_SYSTEM_PROMPT
from schemas.disambiguation_result import DISAMBIGUATION_RESULT_SCHEMA, validate_disambiguation_result

# Import MCP options and server
from openalex_mcp import OPENALEX_MCP_OPTIONS

load_dotenv()

# Global semaphore for concurrency control
# Prevents overwhelming the SDK with too many simultaneous initializations
_DEFAULT_MAX_CONCURRENT = 5
_concurrency_semaphore = None


def _get_semaphore(max_concurrent: int = None) -> asyncio.Semaphore:
    """Get or create the global concurrency semaphore."""
    global _concurrency_semaphore
    if _concurrency_semaphore is None:
        max_concurrent = max_concurrent or _DEFAULT_MAX_CONCURRENT
        _concurrency_semaphore = asyncio.Semaphore(max_concurrent)
    return _concurrency_semaphore


def print_message(message):
    """Print a message from the agent (adapted from agent-ref)."""
    msg_type = type(message).__name__

    # Check if it's an AssistantMessage or UserMessage (they have content attribute)
    if hasattr(message, 'content') and not hasattr(message, 'subtype'):
        # For AssistantMessage and UserMessage, content is a list of blocks
        if message.content:
            for block in message.content:
                block_type = type(block).__name__
                if block_type == 'TextBlock' and hasattr(block, 'text'):
                    print(block.text)
                elif block_type == 'ToolUseBlock':
                    print(f"[Tool Use: {getattr(block, 'name', 'unknown')}]")
                    if hasattr(block, 'id'):
                        print(f"  ID: {block.id}")
                    if hasattr(block, 'input'):
                        input_str = str(block.input)
                        if len(input_str) > 150:
                            input_str = input_str[:147] + "..."
                        print(f"  Input: {input_str}")
                elif block_type == 'ToolResultBlock':
                    if hasattr(block, 'tool_use_id'):
                        print(f"[Tool Result: {block.tool_use_id}]")
                    if hasattr(block, 'is_error') and block.is_error:
                        print(f"  Error: {block.is_error}")
                else:
                    # Fallback: try old attribute names for backwards compatibility
                    if hasattr(block, 'type'):
                        if block.type == "text" and hasattr(block, 'text'):
                            print(block.text)
                        elif block.type == "tool_use":
                            print(f"  🔧 Tool: {block.name}")
        else:
            print(message)
    elif msg_type == 'ResultMessage' or hasattr(message, 'subtype'):
        # ResultMessage doesn't have content, it has other attributes
        print(f"Result: {getattr(message, 'subtype', 'unknown')}")
        if hasattr(message, 'total_cost_usd'):
            print(f"Cost: ${message.total_cost_usd:.6f}")
        if hasattr(message, 'usage'):
            print(f"Usage: {message.usage}")
    else:
        # For other message types, just print the whole message
        print(message)


async def disambiguate_author(
    first_name: str = None,
    last_name: str = None,
    researcher_name: str = None,
    affiliations: list = None,
    institution: str = None,
    context: str = None,
    orcid: str = None,
    find_email: bool = False,
    max_iterations: int = 15,
    max_concurrent: int = None
) -> dict:
    """
    Disambiguate a researcher and find their OpenAlex author ID.

    Args:
        first_name: First name of the researcher (optional if researcher_name provided)
        last_name: Last name of the researcher (optional if researcher_name provided)
        researcher_name: Full name of the researcher (optional if first_name/last_name provided)
        affiliations: List of institution affiliations (optional)
        institution: Single institution affiliation (optional, for backwards compatibility)
        context: Research area, group, project, abstract, or other context (optional)
        orcid: ORCID identifier if known (optional)
        find_email: Whether to search for email address (default: False)
        max_iterations: Maximum iterations (default: 15)
        max_concurrent: Maximum concurrent SDK initializations (default: 5). 
                       Set to None to disable concurrency control. 
                       Lower values prevent timeout errors when processing many authors.

    Returns:
        Dict with disambiguation results (JSON format)

    Note:
        This function uses the Claude Agent SDK with embedded MCP pattern.
        MCP tools are provided via create_sdk_mcp_server() and called by the agent.
        
        Concurrency Control:
        - By default, limits to 5 simultaneous SDK initializations to prevent timeouts
        - If calling disambiguate_author() many times concurrently, keep max_concurrent=5
        - For single calls or low concurrency, set max_concurrent=None for faster execution
    """

    # Determine name - prefer first_name/last_name, fallback to researcher_name
    if first_name and last_name:
        full_name = f"{first_name} {last_name}"
    elif researcher_name:
        full_name = researcher_name
    else:
        raise ValueError("Either (first_name and last_name) or researcher_name must be provided")

    # Handle affiliations - prefer affiliations list, fallback to institution
    institution_list = []
    if affiliations:
        institution_list = affiliations
    elif institution:
        institution_list = [institution]

    # Build user message according to system prompt format
    user_message = "Find the OpenAlex author ID for:\n\n"
    
    if first_name:
        user_message += f"First name: {first_name}\n"
    if last_name:
        user_message += f"Last name: {last_name}\n"
    if not (first_name and last_name) and researcher_name:
        # If only full name provided, try to infer
        name_parts = researcher_name.strip().split(maxsplit=1)
        if len(name_parts) >= 1:
            user_message += f"First name: {name_parts[0]}\n"
        if len(name_parts) >= 2:
            user_message += f"Last name: {name_parts[1]}\n"
    
    if institution_list:
        # Use the first institution as primary, list others if multiple
        user_message += f"institution: {institution_list[0]}\n"
        if len(institution_list) > 1:
            user_message += f"\nAdditional affiliations:\n"
            for aff in institution_list[1:]:
                user_message += f"- {aff}\n"
    
    if context:
        user_message += f"context: {context}\n"
    
    if orcid:
        user_message += f"ORCID: {orcid}\n"
    
    # Add email search parameter
    user_message += f"find_email: {str(find_email).lower()}\n"

    print(f"\n{'='*80}")
    print(f"AUTHOR DISAMBIGUATION REQUEST (Claude SDK + MCP)")
    print(f"{'='*80}")
    print(f"Name: {full_name}")
    if institution_list:
        print(f"Institution(s): {', '.join(institution_list)}")
    if context:
        print(f"Context: {context}")
    if orcid:
        print(f"ORCID: {orcid}")
    print(f"Find email: {find_email}")
    print(f"{'='*80}\n")

    # Configure Claude Agent options with OpenAlex MCP server
    # Using direct prompt with explicit format instructions (no skills needed)
    agent_options = ClaudeAgentOptions(
        mcp_servers=OPENALEX_MCP_OPTIONS.mcp_servers,
        allowed_tools=OPENALEX_MCP_OPTIONS.allowed_tools,
        system_prompt=DIRECT_SYSTEM_PROMPT,  # Explicit instructions without skills
        output_format=DISAMBIGUATION_RESULT_SCHEMA,  # Enforce structured output
    )

    # Track stats
    stats = {
        "input_tokens": 0,
        "output_tokens": 0,
        "iterations": 0,
    }

    final_json = None

    try:
        # Use semaphore for concurrency control if specified
        semaphore = _get_semaphore(max_concurrent) if max_concurrent is not None else None
        
        # Use ClaudeSDKClient with async context manager
        # Wrap with semaphore to prevent too many simultaneous initializations
        async def _run_with_client():
            async with ClaudeSDKClient(options=agent_options) as client:
                # Send the query
                await client.query(user_message)

                # Receive and process response
                all_text = ""
                async for message in client.receive_response():
                    stats["iterations"] += 1

                    # Print the message for user feedback
                    print_message(message)

                    # Collect text for JSON extraction - handle both TextBlock and old-style text
                    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

                    # Track usage if available
                    if hasattr(message, 'usage'):
                        if hasattr(message.usage, 'input_tokens'):
                            stats["input_tokens"] += message.usage.input_tokens
                        if hasattr(message.usage, 'output_tokens'):
                            stats["output_tokens"] += message.usage.output_tokens

                # 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:
                    parsed_json = json.loads(text)
                    return parsed_json
                except json.JSONDecodeError:
                    # Try to extract JSON with regex
                    import re
                    json_match = re.search(r'\{.*\}', text, re.DOTALL)
                    if json_match:
                        try:
                            return json.loads(json_match.group())
                        except:
                            pass
                return None
        
        # Execute with or without semaphore
        if semaphore:
            async with semaphore:
                final_json = await _run_with_client()
        else:
            final_json = await _run_with_client()
        
        # Validate against schema if we got JSON
        if final_json:
            is_valid, validation_errors = validate_disambiguation_result(final_json)
            if not is_valid:
                print(f"\n⚠️  Schema validation warnings:")
                for error in validation_errors:
                    print(f"  - {error}")
                # Still return the result, but log the validation issues

    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return {
            "status": "error",
            "author_candidates": [],
            "search_summary": {
                "search_strategy": "error",
                "tools_used": [],
                "total_candidates_found": 0
            },
            "comments": f"Error during disambiguation: {str(e)}",
            "_metadata": {
                "iterations": stats["iterations"],
                "stats": stats,
                "researcher_name": full_name
            }
        }

    # If we couldn't parse JSON, create error response
    if not final_json:
        final_json = {
            "status": "error",
            "author_candidates": [],
            "search_summary": {
                "search_strategy": "error",
                "tools_used": [],
                "total_candidates_found": 0
            },
            "comments": f"Could not parse JSON response. Raw text: {all_text[:500]}",
            "_metadata": {
                "iterations": stats["iterations"],
                "stats": stats,
                "researcher_name": full_name
            }
        }
    else:
        # Add metadata to successful response
        final_json["_metadata"] = {
            "iterations": stats["iterations"],
            "stats": stats,
            "researcher_name": full_name
        }

    return final_json


def main():
    """CLI entry point."""

    if len(sys.argv) < 2:
        print('Usage: python src/production_agent.py [OPTIONS]')
        print("\nRequired (one of):")
        print('  --first-name "Name"     First name of the author')
        print('  --last-name "Name"     Last name of the author')
        print('  OR')
        print('  --name "Full Name"     Full name of the author (alternative to first/last)')
        print("\nOptions:")
        print('  --affiliation "Name"   Institution affiliation (can be used multiple times)')
        print('  --institution "Name"    Single institution (alternative to --affiliation)')
        print('  --context "Text"       Research area, group, project, abstract, etc.')
        print('  --orcid "ID"           ORCID identifier')
        print('  --find-email           Search for email address')
        print("\nExamples:")
        print('  python src/production_agent.py --first-name "Jerry" --last-name "Adams"')
        print('  python src/production_agent.py --name "Jerry M. Adams" --affiliation "WEHI"')
        print('  python src/production_agent.py --first-name "Yves" --last-name "Barde" --orcid "0000-0002-7627-461X" --find-email')
        print('  python src/production_agent.py --name "John Smith" --affiliation "Harvard" --affiliation "MIT" --context "machine learning"')
        sys.exit(1)

    first_name = None
    last_name = None
    researcher_name = None
    affiliations = []
    institution = None
    context = None
    orcid = None
    find_email = False

    # Parse arguments
    i = 1
    while i < len(sys.argv):
        if sys.argv[i] == '--first-name' and i + 1 < len(sys.argv):
            first_name = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--last-name' and i + 1 < len(sys.argv):
            last_name = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--name' and i + 1 < len(sys.argv):
            researcher_name = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--affiliation' and i + 1 < len(sys.argv):
            affiliations.append(sys.argv[i + 1])
            i += 2
        elif sys.argv[i] == '--institution' and i + 1 < len(sys.argv):
            institution = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--context' and i + 1 < len(sys.argv):
            context = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--orcid' and i + 1 < len(sys.argv):
            orcid = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--find-email':
            find_email = True
            i += 1
        else:
            i += 1

    # Validate that we have name information
    if not (first_name and last_name) and not researcher_name:
        print("Error: Must provide either (--first-name and --last-name) or --name")
        sys.exit(1)

    try:
        # Run the async function
        result = asyncio.run(disambiguate_author(
            first_name=first_name,
            last_name=last_name,
            researcher_name=researcher_name,
            affiliations=affiliations if affiliations else None,
            institution=institution,
            context=context,
            orcid=orcid,
            find_email=find_email
        ))

        print(f"\n{'='*80}")
        print("FINAL RESULT (Claude SDK + MCP)")
        print(f"{'='*80}\n")
        print(json.dumps(result, indent=2))
        print(f"\n{'='*80}\n")

        # Save to file
        name_for_file = (f"{first_name}_{last_name}" if (first_name and last_name) else researcher_name).replace(' ', '_')
        output_file = f"result_{name_for_file}.json"
        with open(output_file, 'w') as f:
            json.dump(result, f, indent=2)

        print(f"💾 Saved to: {output_file}\n")

    except KeyboardInterrupt:
        print("\n\n⚠️  Interrupted")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()
