#!/usr/bin/env python3
"""
README/Markdown Summarization Script
Scans README files in each repository, uses LLM to generate summaries, writes back to README_SUMMARY.md
"""

import os
import sys
import json
import asyncio
import argparse
from pathlib import Path
from typing import List, Dict, Optional
from tqdm import tqdm
import hashlib
from dotenv import load_dotenv

# Load .env file (before importing other modules)
env_file = Path(__file__).parent / ".env"
if env_file.exists():
    load_dotenv(env_file)
elif (Path(__file__).parent.parent / ".env").exists():
    # If not in current directory, try loading from project root
    load_dotenv(Path(__file__).parent.parent / ".env")

# Add current directory to path (for importing schemas)
sys.path.insert(0, str(Path(__file__).parent))
# Add domain_code/src to path for reusing util functions
sys.path.insert(0, str(Path(__file__).parent.parent / "domain_code" / "src"))
from util import call_llm, init_logger, logger
from schemas import READMESummary


# Default output filename (written back to repository directory)
SUMMARY_FILENAME = "README_SUMMARY.md"


def find_readme_files(repo_dir: Path) -> List[Path]:
    """
    Find README files in the repository
    
    Args:
        repo_dir: Repository root directory
        
    Returns:
        List of README file paths
    """
    readme_files = []
    
    # README* files in repository root directory
    for pattern in ["README*", "readme*"]:
        readme_files.extend(repo_dir.glob(pattern))
    
    # README* files in docs/ directory
    docs_dir = repo_dir / "docs"
    if docs_dir.exists() and docs_dir.is_dir():
        for pattern in ["README*", "readme*"]:
            readme_files.extend(docs_dir.glob(pattern))
    
    # Filter: keep only .md or .markdown files
    readme_files = [
        f for f in readme_files 
        if f.is_file() and f.suffix.lower() in [".md", ".markdown"]
    ]
    
    return sorted(set(readme_files))  # Deduplicate and sort


def read_readme_content(readme_files: List[Path]) -> str:
    """
    Read and merge all README file contents
    
    Args:
        readme_files: List of README file paths
        
    Returns:
        Merged README content
    """
    contents = []
    for readme_file in readme_files:
        try:
            with open(readme_file, "r", encoding="utf-8", errors="ignore") as f:
                content = f.read().strip()
                if content:
                    contents.append(f"## File: {readme_file.name}\n\n{content}")
        except Exception as e:
            logger.warning(f"Unable to read file {readme_file}: {e}")
    
    return "\n\n---\n\n".join(contents)


async def summarize_readme(
    readme_content: str,
    base_url: str,
    model: str,
    api_key: str,
    log_file: str,
) -> Optional[Dict]:
    """
    Use LLM to summarize README content
    
    Args:
        readme_content: README file content
        base_url: LLM API base URL
        model: Model name
        api_key: API key
        log_file: Log file path
        
    Returns:
        README summary (dict) or None
    """
    # Read prompt template
    prompt_template_path = Path(__file__).parent / "prompts" / "readme_summary.txt"
    try:
        with open(prompt_template_path, "r", encoding="utf-8") as f:
            prompt_template = f.read()
    except Exception as e:
        logger.error(f"Unable to read prompt template: {e}")
        return None
    
    # Build prompt
    prompt = prompt_template.format(readme_content=readme_content)
    
    # Call LLM
    messages = [{"role": "user", "content": prompt}]
    
    try:
        result = await call_llm(
            messages=messages,
            model=model,
            base_url=base_url,
            api_key=api_key,
            pydantic_object=READMESummary,
            log_file=log_file,
        )
        
        if result is None:
            logger.warning("LLM call returned None, skipping")
            return None
        
        # If result is a string, try to parse JSON
        if isinstance(result, str):
            try:
                result = json.loads(result)
            except json.JSONDecodeError:
                logger.warning(f"Unable to parse JSON from LLM response: {result[:200]}")
                return None
        
        return result
    except Exception as e:
        logger.error(f"LLM call failed: {e}")
        return None


def write_summary_file(
    repo_dir: Path,
    summary: Dict,
    readme_files: List[Path],
) -> Path:
    """
    Write summary to README_SUMMARY.md file (in repository directory)
    
    Args:
        repo_dir: Repository root directory
        summary: README summary (dict)
        readme_files: Original README file list
        
    Returns:
        Output file path
    """
    output_file = repo_dir / SUMMARY_FILENAME
    
    # Build output content (simplified format)
    lines = []
    lines.append("# Project Summary\n\n")
    
    # Add summary fields (only essential ones)
    if "project_overview" in summary:
        lines.append(f"## Project Overview\n\n{summary['project_overview']}\n\n")
    
    if "main_features" in summary:
        lines.append(f"## Main Features\n\n{summary['main_features']}\n\n")
    
    # Write to file
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            f.write("".join(lines))
        logger.info(f"Summary file written: {output_file}")
        return output_file
    except Exception as e:
        logger.error(f"Unable to write summary file {output_file}: {e}")
        raise


async def process_single_repo(
    repo_dir: Path,
    base_url: str,
    model: str,
    api_key: str,
    log_file: str,
    overwrite: bool = False,
) -> Dict[str, any]:
    """
    Process README summarization for a single repository
    
    Args:
        repo_dir: Repository root directory
        base_url: LLM API base URL
        model: Model name
        api_key: API key
        log_file: Log file path
        overwrite: Whether to overwrite existing summary file
        
    Returns:
        Processing result dictionary
    """
    repo_name = repo_dir.name
    summary_file = repo_dir / SUMMARY_FILENAME
    
    # Check if summary file already exists
    if summary_file.exists() and not overwrite:
        return {
            "repo": repo_name,
            "status": "skipped",
            "reason": "Summary file already exists",
        }
    
    # Find README files
    readme_files = find_readme_files(repo_dir)
    if not readme_files:
        return {
            "repo": repo_name,
            "status": "no_readme",
            "reason": "README file not found",
        }
    
    # Read README content
    readme_content = read_readme_content(readme_files)
    if not readme_content:
        return {
            "repo": repo_name,
            "status": "empty_readme",
            "reason": "README file is empty",
        }
    
    # Call LLM to generate summary
    summary = await summarize_readme(
        readme_content=readme_content,
        base_url=base_url,
        model=model,
        api_key=api_key,
        log_file=log_file,
    )
    
    if summary is None:
        return {
            "repo": repo_name,
            "status": "llm_failed",
            "reason": "LLM call failed",
        }
    
    # Write summary file
    try:
        write_summary_file(repo_dir, summary, readme_files)
        return {
            "repo": repo_name,
            "status": "success",
            "summary_file": str(summary_file),
            "readme_count": len(readme_files),
        }
    except Exception as e:
        return {
            "repo": repo_name,
            "status": "write_failed",
            "reason": str(e),
        }


async def process_all_repos(
    repos_dir: Path,
    base_url: str,
    model: str,
    api_key: str,
    log_file: str,
    max_concurrency: int = 8,
    overwrite: bool = False,
) -> List[Dict]:
    """
    Process README summarization for all repositories
    
    Args:
        repos_dir: Repository root directory
        base_url: LLM API base URL
        model: Model name
        api_key: API key
        log_file: Log file path
        max_concurrency: Maximum concurrency
        overwrite: Whether to overwrite existing summary files
        
    Returns:
        List of processing results for all repositories
    """
    # Get all repository directories
    repo_dirs = [
        d for d in repos_dir.iterdir() 
        if d.is_dir() and not d.name.startswith(".")
    ]
    repo_dirs.sort()
    
    logger.info(f"Found {len(repo_dirs)} repositories, starting processing...")
    
    # Use semaphore to control concurrency
    semaphore = asyncio.Semaphore(max_concurrency)
    
    async def process_with_semaphore(repo_dir: Path):
        async with semaphore:
            return await process_single_repo(
                repo_dir=repo_dir,
                base_url=base_url,
                model=model,
                api_key=api_key,
                log_file=log_file,
                overwrite=overwrite,
            )
    
    # Process all repositories concurrently
    tasks = [process_with_semaphore(repo_dir) for repo_dir in repo_dirs]
    results = []
    
    for task in tqdm(asyncio.as_completed(tasks), total=len(tasks), desc="Processing repos"):
        result = await task
        results.append(result)
    
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="README/Markdown Summarization Tool")
    parser.add_argument(
        "--repos_dir",
        type=str,
        default="/home/weifengsun/tangou1/domain_code/src/workdir/repos_filtered",
        help="Repository root directory path",
    )
    parser.add_argument(
        "--base_url",
        type=str,
        default=os.getenv("OPENAI_BASE_URL", "http://localhost:8000/v1"),
        help="LLM API base URL (default: http://localhost:8000/v1)",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="Qwen3",
        help="Model name (default: Qwen3)",
    )
    parser.add_argument(
        "--api_key_env",
        type=str,
        default="OPENAI_API_KEY",
        help="API key environment variable name (default: OPENAI_API_KEY)",
    )
    parser.add_argument(
        "--max_concurrency",
        type=int,
        default=8,
        help="Maximum concurrency (default: 8)",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing summary files",
    )
    parser.add_argument(
        "--log_file",
        type=str,
        default="instruction_generation/workdir/logs/summarize.log",
        help="Log file path",
    )
    
    args = parser.parse_args()
    
    # Initialize logger
    init_logger(args.log_file, level="INFO")
    
    # Get API key
    api_key = os.getenv(args.api_key_env, "none")
    
    # Process all repositories
    repos_dir = Path(args.repos_dir)
    if not repos_dir.exists():
        logger.error(f"Repository directory does not exist: {repos_dir}")
        sys.exit(1)
    
    # Create log directory
    log_file_path = Path(args.log_file)
    log_file_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Run main logic
    results = asyncio.run(
        process_all_repos(
            repos_dir=repos_dir,
            base_url=args.base_url,
            model=args.model,
            api_key=api_key,
            log_file=str(log_file_path),
            max_concurrency=args.max_concurrency,
            overwrite=args.overwrite,
        )
    )
    
    # Statistics
    status_counts = {}
    for result in results:
        status = result["status"]
        status_counts[status] = status_counts.get(status, 0) + 1
    
    logger.info("\n" + "=" * 80)
    logger.info("Processing complete!")
    logger.info("=" * 80)
    logger.info(f"Total: {len(results)} repositories")
    for status, count in status_counts.items():
        logger.info(f"  {status}: {count}")
    logger.info("=" * 80)