#!/usr/bin/env python3
"""
Code Reference Indexer MCP Tool - Unified Version

Specialized MCP tool for searching relevant index content in indexes folder
and formatting it for LLM code implementation reference.

Core Features:
1. **UNIFIED TOOL**: Combined search_code_references that handles directory setup, loading, and searching in one call
2. Match relevant reference code based on target file path and functionality requirements
3. Format output of relevant code examples, functions and concepts
4. Provide structured reference information for LLM use

Key Improvement:
- Single tool call that handles all steps internally
- Agent only needs to provide indexes_path and target_file
- No dependency on calling order or global state management
"""

import json
from pathlib import Path
from typing import Dict, List, Tuple
from dataclasses import dataclass
import logging

# Import MCP modules
from mcp.server.fastmcp import FastMCP

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create FastMCP server instance
mcp = FastMCP("code-reference-indexer")


@dataclass
class CodeReference:
 """Code reference information structure"""

 file_path: str
 file_type: str
 main_functions: List[str]
 key_concepts: List[str]
 dependencies: List[str]
 summary: str
 lines_of_code: int
 repo_name: str
 confidence_score: float = 0.0


@dataclass
class RelationshipInfo:
 """Relationship information structure"""

 repo_file_path: str
 target_file_path: str
 relationship_type: str
 confidence_score: float
 helpful_aspects: List[str]
 potential_contributions: List[str]
 usage_suggestions: str


def load_index_files_from_directory(indexes_directory: str) -> Dict[str, Dict]:
 """Load all index files from specified directory"""
 indexes_path = Path(indexes_directory).resolve()

 if not indexes_path.exists():
 logger.warning(f"Indexes directory does not exist: {indexes_path}")
 return {}

 index_cache = {}

 for index_file in indexes_path.glob("*.json"):
 try:
 with open(index_file, "r", encoding="utf-8") as f:
 index_data = json.load(f)
 index_cache[index_file.stem] = index_data
 logger.info(f"Loaded index file: {index_file.name}")
 except Exception as e:
 logger.error(f"Failed to load index file {index_file.name}: {e}")

 logger.info(f"Loaded {len(index_cache)} index files from {indexes_path}")
 return index_cache


def extract_code_references(index_data: Dict) -> List[CodeReference]:
 """Extract code reference information from index data"""
 references = []

 repo_name = index_data.get("repo_name", "Unknown")
 file_summaries = index_data.get("file_summaries", [])

 for file_summary in file_summaries:
 reference = CodeReference(
 file_path=file_summary.get("file_path", ""),
 file_type=file_summary.get("file_type", ""),
 main_functions=file_summary.get("main_functions", []),
 key_concepts=file_summary.get("key_concepts", []),
 dependencies=file_summary.get("dependencies", []),
 summary=file_summary.get("summary", ""),
 lines_of_code=file_summary.get("lines_of_code", 0),
 repo_name=repo_name,
 )
 references.append(reference)

 return references


def extract_relationships(index_data: Dict) -> List[RelationshipInfo]:
 """Extract relationship information from index data"""
 relationships = []

 relationship_list = index_data.get("relationships", [])

 for rel in relationship_list:
 relationship = RelationshipInfo(
 repo_file_path=rel.get("repo_file_path", ""),
 target_file_path=rel.get("target_file_path", ""),
 relationship_type=rel.get("relationship_type", ""),
 confidence_score=rel.get("confidence_score", 0.0),
 helpful_aspects=rel.get("helpful_aspects", []),
 potential_contributions=rel.get("potential_contributions", []),
 usage_suggestions=rel.get("usage_suggestions", ""),
 )
 relationships.append(relationship)

 return relationships


def calculate_relevance_score(
 target_file: str, reference: CodeReference, keywords: List[str] = None
) -> float:
 """Calculate relevance score between reference code and target file"""
 score = 0.0

 # File name similarity
 target_name = Path(target_file).stem.lower()
 ref_name = Path(reference.file_path).stem.lower()

 if target_name in ref_name or ref_name in target_name:
 score += 0.3

 # File type matching
 target_extension = Path(target_file).suffix
 ref_extension = Path(reference.file_path).suffix

 if target_extension == ref_extension:
 score += 0.2

 # Keyword matching
 if keywords:
 keyword_matches = 0
 total_searchable_text = (
 " ".join(reference.key_concepts)
 + " "
 + " ".join(reference.main_functions)
 + " "
 + reference.summary
 + " "
 + reference.file_type
 ).lower()

 for keyword in keywords:
 if keyword.lower() in total_searchable_text:
 keyword_matches += 1

 if keywords:
 score += (keyword_matches / len(keywords)) * 0.5

 return min(score, 1.0)


def find_relevant_references_in_cache(
 target_file: str,
 index_cache: Dict[str, Dict],
 keywords: List[str] = None,
 max_results: int = 10,
) -> List[Tuple[CodeReference, float]]:
 """Find reference code relevant to target file from provided cache"""
 all_references = []

 # Collect reference information from all index files
 for repo_name, index_data in index_cache.items():
 references = extract_code_references(index_data)
 for ref in references:
 relevance_score = calculate_relevance_score(target_file, ref, keywords)
 if relevance_score > 0.1: # Only keep results with certain relevance
 all_references.append((ref, relevance_score))

 # Sort by relevance score
 all_references.sort(key=lambda x: x[1], reverse=True)

 return all_references[:max_results]


def find_direct_relationships_in_cache(
 target_file: str, index_cache: Dict[str, Dict]
) -> List[RelationshipInfo]:
 """Find direct relationships with target file from provided cache"""
 relationships = []

 # Normalize target file path (remove common prefixes if exists)
 common_prefixes = ["src/", "core/", "lib/", "main/", "./"]
 normalized_target = target_file.strip("/")
 for prefix in common_prefixes:
 if normalized_target.startswith(prefix):
 normalized_target = normalized_target[len(prefix) :]
 break

 # Collect relationship information from all index files
 for repo_name, index_data in index_cache.items():
 repo_relationships = extract_relationships(index_data)
 for rel in repo_relationships:
 # Normalize target file path in relationship
 normalized_rel_target = rel.target_file_path.strip("/")
 for prefix in common_prefixes:
 if normalized_rel_target.startswith(prefix):
 normalized_rel_target = normalized_rel_target[len(prefix) :]
 break

 # Check target file path matching (support multiple matching methods)
 if (
 normalized_target == normalized_rel_target
 or normalized_target in normalized_rel_target
 or normalized_rel_target in normalized_target
 or target_file in rel.target_file_path
 or rel.target_file_path in target_file
 ):
 relationships.append(rel)

 # Sort by confidence score
 relationships.sort(key=lambda x: x.confidence_score, reverse=True)

 return relationships


def format_reference_output(
 target_file: str,
 relevant_refs: List[Tuple[CodeReference, float]],
 relationships: List[RelationshipInfo],
) -> str:
 """Format reference information output"""
 output_lines = []

 output_lines.append(f"# Code Reference Information - {target_file}")
 output_lines.append("=" * 80)
 output_lines.append("")

 # Direct relationship information
 if relationships:
 output_lines.append("## 🎯 Direct Relationships")
 output_lines.append("")

 for i, rel in enumerate(relationships[:5], 1):
 output_lines.append(f"### {i}. {rel.repo_file_path}")
 output_lines.append(f"**Relationship Type**: {rel.relationship_type}")
 output_lines.append(f"**Confidence Score**: {rel.confidence_score:.2f}")
 output_lines.append(
 f"**Helpful Aspects**: {', '.join(rel.helpful_aspects)}"
 )
 output_lines.append(
 f"**Potential Contributions**: {', '.join(rel.potential_contributions)}"
 )
 output_lines.append(f"**Usage Suggestions**: {rel.usage_suggestions}")
 output_lines.append("")

 # Relevant code references
 if relevant_refs:
 output_lines.append("## 📚 Relevant Code References")
 output_lines.append("")

 for i, (ref, score) in enumerate(relevant_refs[:8], 1):
 output_lines.append(f"### {i}. {ref.file_path} (Relevance: {score:.2f})")
 output_lines.append(f"**Repository**: {ref.repo_name}")
 output_lines.append(f"**File Type**: {ref.file_type}")
 output_lines.append(
 f"**Main Functions**: {', '.join(ref.main_functions[:5])}"
 )
 output_lines.append(f"**Key Concepts**: {', '.join(ref.key_concepts[:8])}")
 output_lines.append(f"**Dependencies**: {', '.join(ref.dependencies[:6])}")
 output_lines.append(f"**Lines of Code**: {ref.lines_of_code}")
 output_lines.append(f"**Summary**: {ref.summary[:300]}...")
 output_lines.append("")

 # Implementation suggestions
 output_lines.append("## 💡 Implementation Suggestions")
 output_lines.append("")

 if relevant_refs:
 # Collect all function names and concepts
 all_functions = set()
 all_concepts = set()
 all_dependencies = set()

 for ref, _ in relevant_refs[:5]:
 all_functions.update(ref.main_functions)
 all_concepts.update(ref.key_concepts)
 all_dependencies.update(ref.dependencies)

 output_lines.append("**Reference Function Name Patterns**:")
 for func in sorted(list(all_functions))[:10]:
 output_lines.append(f"- {func}")
 output_lines.append("")

 output_lines.append("**Important Concepts and Patterns**:")
 for concept in sorted(list(all_concepts))[:15]:
 output_lines.append(f"- {concept}")
 output_lines.append("")

 output_lines.append("**Potential Dependencies Needed**:")
 for dep in sorted(list(all_dependencies))[:10]:
 output_lines.append(f"- {dep}")
 output_lines.append("")

 output_lines.append("## 🚀 Next Actions")
 output_lines.append(
 "1. Analyze design patterns and architectural styles from the above reference code"
 )
 output_lines.append("2. Determine core functionalities and interfaces to implement")
 output_lines.append("3. Choose appropriate dependency libraries and tools")
 output_lines.append(
 "4. Design implementation solution consistent with existing code style"
 )
 output_lines.append("5. Start writing specific code implementation")

 return "\n".join(output_lines)


# ==================== MCP Tool Definitions ====================


@mcp.tool()
async def search_code_references(
 indexes_path: str, target_file: str, keywords: str = "", max_results: int = 10
) -> str:
 """
 **UNIFIED TOOL**: Search relevant reference code from index files for target file implementation.
 This tool combines directory setup, index loading, and searching in a single call.

 Args:
 indexes_path: Path to the indexes directory containing JSON index files
 target_file: Target file path (file to be implemented)
 keywords: Search keywords, comma-separated
 max_results: Maximum number of results to return

 Returns:
 Formatted reference code information JSON string
 """
 try:
 # Step 1: Load index files from specified directory
 logger.info(f"Loading index files from: {indexes_path}")
 index_cache = load_index_files_from_directory(indexes_path)

 if not index_cache:
 result = {
 "status": "error",
 "message": f"No index files found or failed to load from: {indexes_path}",
 "target_file": target_file,
 "indexes_path": indexes_path,
 }
 return json.dumps(result, ensure_ascii=False, indent=2)

 # Step 2: Parse keywords
 keyword_list = (
 [kw.strip() for kw in keywords.split(",") if kw.strip()] if keywords else []
 )

 # Step 3: Find relevant reference code
 relevant_refs = find_relevant_references_in_cache(
 target_file, index_cache, keyword_list, max_results
 )

 # Step 4: Find direct relationships
 relationships = find_direct_relationships_in_cache(target_file, index_cache)

 # Step 5: Format output
 formatted_output = format_reference_output(
 target_file, relevant_refs, relationships
 )

 result = {
 "status": "success",
 "target_file": target_file,
 "indexes_path": indexes_path,
 "keywords_used": keyword_list,
 "total_references_found": len(relevant_refs),
 "total_relationships_found": len(relationships),
 "formatted_content": formatted_output,
 "indexes_loaded": list(index_cache.keys()),
 "total_indexes_loaded": len(index_cache),
 }

 logger.info(
 f"Successfully found {len(relevant_refs)} references and {len(relationships)} relationships for {target_file}"
 )
 return json.dumps(result, ensure_ascii=False, indent=2)

 except Exception as e:
 logger.error(f"Error in search_code_references: {str(e)}")
 result = {
 "status": "error",
 "message": f"Failed to search reference code: {str(e)}",
 "target_file": target_file,
 "indexes_path": indexes_path,
 }
 return json.dumps(result, ensure_ascii=False, indent=2)


@mcp.tool()
async def get_indexes_overview(indexes_path: str) -> str:
 """
 Get overview of all available reference code index information from specified directory

 Args:
 indexes_path: Path to the indexes directory containing JSON index files

 Returns:
 Overview information of all available reference code JSON string
 """
 try:
 # Load index files from specified directory
 index_cache = load_index_files_from_directory(indexes_path)

 if not index_cache:
 result = {
 "status": "error",
 "message": f"No index files found in: {indexes_path}",
 "indexes_path": indexes_path,
 }
 return json.dumps(result, ensure_ascii=False, indent=2)

 overview = {"total_repos": len(index_cache), "repositories": {}}

 for repo_name, index_data in index_cache.items():
 repo_info = {
 "repo_name": index_data.get("repo_name", repo_name),
 "total_files": index_data.get("total_files", 0),
 "file_types": [],
 "main_concepts": [],
 "total_relationships": len(index_data.get("relationships", [])),
 }

 # Collect file types and concepts
 file_summaries = index_data.get("file_summaries", [])
 file_types = set()
 concepts = set()

 for file_summary in file_summaries:
 file_types.add(file_summary.get("file_type", "Unknown"))
 concepts.update(file_summary.get("key_concepts", []))

 repo_info["file_types"] = sorted(list(file_types))
 repo_info["main_concepts"] = sorted(list(concepts))[
 :20
 ] # Limit concept count

 overview["repositories"][repo_name] = repo_info

 result = {
 "status": "success",
 "overview": overview,
 "indexes_directory": str(Path(indexes_path).resolve()),
 "total_indexes_loaded": len(index_cache),
 }

 return json.dumps(result, ensure_ascii=False, indent=2)

 except Exception as e:
 result = {
 "status": "error",
 "message": f"Failed to get indexes overview: {str(e)}",
 "indexes_path": indexes_path,
 }
 return json.dumps(result, ensure_ascii=False, indent=2)


def main():
 """Main function"""
 logger.info("Starting unified Code Reference Indexer MCP server")
 logger.info("Available tools:")
 logger.info(
 "1. search_code_references(indexes_path, target_file, keywords, max_results) - UNIFIED TOOL"
 )
 logger.info(
 "2. get_indexes_overview(indexes_path) - Get overview of available indexes"
 )

 # Run MCP server
 mcp.run()


if __name__ == "__main__":
 main()