"""Agent_BRICK query module - HTTP API client for BRICK KG API service.

This module provides query_cypher and query_relation functions that call
the remote BRICK KG API service instead of direct Neo4j connection.
"""

import argparse
import json
import re
import sys
from typing import List, Optional, Union
import pandas as pd
import requests

# BRICK KG API base URL
API_BASE_URL = "https://biomics.dcs.cloud/api"


def _extract_entities_from_text(text: str) -> List[str]:
    """
    Extract entity names from text, handling various formats.
    
    Supports:
    - Comma-separated: "GeneA, GeneB, GeneC"
    - With parentheses: "GeneA (protein), GeneB (gene)"
    - With descriptions: "GeneA - kinase, GeneB"
    - One per line
    - Mixed formats
    
    Args:
        text: Raw text containing entity names.
        
    Returns:
        List of cleaned entity names.
    """
    if not text or not isinstance(text, str):
        return []
    
    # Split by common delimiters (comma, newline, tab)
    # First try to split by newline for one-per-line format
    lines = text.replace('\r\n', '\n').replace('\r', '\n').split('\n')
    
    entities = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        
        # Split line by comma if it contains commas
        parts = [p.strip() for p in line.split(',')]
        
        for part in parts:
            if not part:
                continue
            
            # Extract entity name by removing common annotations
            # Pattern 1: "Entity (description)" -> "Entity"
            # Pattern 2: "Entity - description" -> "Entity"
            # Pattern 3: "Entity [id]" -> "Entity"
            # Pattern 4: "1. Entity" -> "Entity" (numbered list)
            
            # Remove leading numbers (e.g., "1. GeneA" or "1) GeneA")
            part = re.sub(r'^\s*[\d]+[.\)]\s*', '', part)
            
            # Remove content in parentheses, brackets, braces
            part = re.sub(r'\s*[\(\[\{].*?[\)\]\}]', '', part)
            
            # Remove content after common separators
            part = re.sub(r'\s*[\-–—:;].*$', '', part)
            
            # Clean up whitespace
            entity = part.strip()
            
            # Skip if empty or too short (likely not a valid entity)
            if len(entity) < 2:
                continue
                
            entities.append(entity)
    
    # Remove duplicates while preserving order
    seen = set()
    unique_entities = []
    for e in entities:
        key = e.lower()
        if key not in seen:
            seen.add(key)
            unique_entities.append(e)
    
    return unique_entities


def _load_entities_from_file(filepath: str) -> List[str]:
    """
    Load and extract entity names from a file.
    
    Args:
        filepath: Path to the file containing entity names.
        
    Returns:
        List of cleaned entity names.
        
    Raises:
        FileNotFoundError: If file doesn't exist.
        ValueError: If no valid entities found.
    """
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
    except UnicodeDecodeError:
        # Try with different encoding
        with open(filepath, 'r', encoding='latin-1') as f:
            content = f.read()
    
    entities = _extract_entities_from_text(content)
    
    if not entities:
        raise ValueError(f"No valid entities found in file: {filepath}")
    
    return entities


def query_cypher(cypher: str, parameters: Optional[dict] = None) -> pd.DataFrame:
    """
    Execute a Cypher query via BRICK KG API.
    
    Args:
        cypher: The Cypher query string.
        parameters: Query parameters dictionary.
        
    Returns:
        pandas.DataFrame: Query results with flattened columns.
    """
    url = f"{API_BASE_URL}/v1/cypher"
    payload = {
        "cypher": cypher,
        "parameters": parameters or {},
        "return_type": "dataframe"
    }
    
    response = requests.post(url, json=payload, timeout=60)
    response.raise_for_status()
    
    result = response.json()
    if not result.get("success"):
        raise RuntimeError(f"API error: {result.get('message', 'Unknown error')}")
    
    data = result.get("data", {})
    rows = data.get("rows", [])
    columns = data.get("columns", [])
    
    if not rows:
        return pd.DataFrame(columns=columns)
    
    return pd.DataFrame(rows, columns=columns)


def _cli_main():
    """CLI entry point for BRICK_qr command."""
    parser = argparse.ArgumentParser(
        description="BRICK_qr - Query entity relations",
        usage='BRICK_qr "entity_name" [options]'
    )
    
    parser.add_argument("source", nargs="?", help="Source entity name(s), comma-separated for multiple")
    parser.add_argument("-i", "--input-file", help="File containing source entity names (one per line or comma-separated)")
    parser.add_argument("-t", "--target", help="Target entity name(s)")
    parser.add_argument("--target-file", help="File containing target entity names (one per line or comma-separated)")
    parser.add_argument("-r", "--relation", help="Relation type (e.g., marker_of, associated_with)")
    parser.add_argument("--st", help="Source entity type (e.g., Cell, Gene, Disease)")
    parser.add_argument("--tt", help="Target entity type")
    parser.add_argument("-n", type=int, default=10, help="Number of results (default: 10)")
    parser.add_argument("-o", help="Output CSV file")
    
    args = parser.parse_args()
    
    # Validate input: either source or input-file must be provided
    if not args.source and not args.input_file:
        parser.error("Either 'source' argument or '-i/--input-file' must be provided")
    
    try:
        # Load source entities
        if args.input_file:
            print(f"Loading source entities from: {args.input_file}")
            source_entities = _load_entities_from_file(args.input_file)
            print(f"Extracted {len(source_entities)} source entities: {source_entities[:5]}{'...' if len(source_entities) > 5 else ''}")
        else:
            source_entities = _extract_entities_from_text(args.source)
            if not source_entities:
                raise ValueError("No valid source entities found")
            print(f"Source entities: {source_entities}")
        
        # Load target entities
        target_entities = None
        if args.target_file:
            print(f"Loading target entities from: {args.target_file}")
            target_entities = _load_entities_from_file(args.target_file)
            print(f"Extracted {len(target_entities)} target entities: {target_entities[:5]}{'...' if len(target_entities) > 5 else ''}")
        elif args.target:
            target_entities = _extract_entities_from_text(args.target)
            if target_entities:
                print(f"Target entities: {target_entities}")
        
        relations = [r.strip() for r in args.relation.split(",")] if args.relation else None
        
        print(f"\nQuerying {len(source_entities)} source entities...")
        
        df = query_relation(
            source_entity_set=source_entities,
            target_entity_set=target_entities,
            relation=relations,
            source_entity_type=args.st,
            target_entity_type=args.tt,
            limit_number=args.n
        )
        
        if df.empty:
            print("No results found.")
            return
        
        print(f"Found {len(df)} relations")
        
        if args.o:
            df.to_csv(args.o, index=False)
            print(f"Saved to {args.o}")
        else:
            cols = ['path.0.name', 'path.1.relation', 'path.2.name', 'path.2.type']
            available = [c for c in cols if c in df.columns]
            print(df[available].to_string())
            
    except FileNotFoundError as e:
        print(f"Error: File not found - {e}", file=sys.stderr)
        sys.exit(1)
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    _cli_main()


def query_relation(
    source_entity_set: Union[str, List[str]],
    target_entity_set: Optional[Union[str, List[str]]] = None,
    relation: Optional[Union[str, List[str]]] = None,
    multi_hop: Union[int, tuple] = 1,
    directed: bool = False,
    source_entity_type: Optional[str] = None,
    target_entity_type: Optional[str] = None,
    relation_info_source: Optional[Union[str, List[str]]] = None,
    query_attribution: str = "name",
    limit_number: Optional[int] = None,
    return_type: str = "dataframe"
) -> pd.DataFrame:
    """
    Query relationships between entities via BRICK KG API.
    
    Supports two modes:
    - Neighbor mode (target_entity_set=None): finds all connected entities.
    - Relation mode (target_entity_set provided): validates specific pairs.
    
    Args:
        source_entity_set: Source entity or list of entities.
        target_entity_set: Target entity or list (None for neighbor mode).
        relation: Relationship type(s) to filter.
        multi_hop: Number of hops (int or tuple for range).
        directed: Whether to query directed paths.
        source_entity_type: Type of source entities (e.g., 'Cell', 'Gene').
        target_entity_type: Type of target entities.
        relation_info_source: Filter by info source.
        query_attribution: Attribute to match ('name' or 'id').
        limit_number: Maximum results to return.
        return_type: Return format ('dataframe' or 'graph').
        
    Returns:
        pandas.DataFrame: Query results with relationship data.
    """
    url = f"{API_BASE_URL}/v1/relation"
    
    # Normalize source_entity_set to list
    if isinstance(source_entity_set, str):
        source_entity_set = [source_entity_set]
    else:
        source_entity_set = list(source_entity_set)
    
    # Normalize target_entity_set
    if target_entity_set is not None:
        if isinstance(target_entity_set, str):
            target_entity_set = [target_entity_set]
        else:
            target_entity_set = list(target_entity_set)
    
    # Normalize relation to list if provided
    if relation is not None and isinstance(relation, str):
        relation = [relation]
    
    # Normalize multi_hop
    if isinstance(multi_hop, tuple):
        multi_hop = list(multi_hop)
    
    payload = {
        "source_entity_set": source_entity_set,
        "target_entity_set": target_entity_set,
        "relation": relation,
        "multi_hop": multi_hop,
        "directed": directed,
        "source_entity_type": source_entity_type,
        "target_entity_type": target_entity_type,
        "relation_info_source": relation_info_source,
        "query_attribution": query_attribution,
        "limit_number": limit_number,
        "return_type": return_type
    }
    
    # Remove None values
    payload = {k: v for k, v in payload.items() if v is not None}
    
    response = requests.post(url, json=payload, timeout=60)
    response.raise_for_status()
    
    result = response.json()
    if not result.get("success"):
        raise RuntimeError(f"API error: {result.get('message', 'Unknown error')}")
    
    data = result.get("data", {})
    rows = data.get("rows", [])
    columns = data.get("columns", [])
    
    if not rows:
        return pd.DataFrame(columns=columns)
    
    return pd.DataFrame(rows, columns=columns)
