import os
import re
import sys
import json
from pathlib import Path
import shutil

filenames = [
    "docs/changes.md",
    "docs/overview.md",
    "docs/lexical_structure.md",
    "docs/primitive_types.md",
    "docs/basic_operations.md",
    "docs/functions.md",
    "docs/control_flow.md",
    "docs/data_structures.md",
    "docs/type_system.md",
    "docs/trait_system.md",
    "docs/generic_programming.md",
    "docs/advanced_topics.md",
    "docs/standard_library.md",
    "docs/complete_trait_reference.md",
    "docs/complete_operator_reference.md",
    "docs/grammar_summary.md",
    "docs/module_system.md",
    "docs/memory_management.md",
    "docs/examples_and_tutorials.md",
    "docs/complete_type_reference.md",
]

def remove_front_matter(content):
    """Remove YAML front matter (--- ... ---) from markdown content."""
    pattern = r'^---\n.*?\n---\n'
    return re.sub(pattern, '', content, flags=re.DOTALL)

def extract_front_matter(content):
    """Extract YAML front matter from markdown content."""
    match = re.match(r'^---\n(.*?)\n---\n', content, flags=re.DOTALL)
    return match.group(1).strip() if match else None

def read_file_list(input_arg):
    """Return list of files either from arguments or a text file of paths."""
    p = Path(input_arg)
    if p.is_file() and p.suffix == '.txt':
        with p.open(encoding='utf-8') as f:
            return [line.strip() for line in f if line.strip()]
    return [input_arg]

def combine_markdown(file_inputs, output_combined, output_meta_json):
    """Combine multiple markdown files into one, skipping front matter."""
    files = []
    for f in file_inputs:
        files.extend(read_file_list(f))

    combined_parts = []
    meta_info = {"order": [], "files": {}}

    for file in files:
        file_path = Path(file)
        if not file_path.exists():
            print(f"Skipping missing file: {file}")
            continue

        text = file_path.read_text(encoding='utf-8')
        front_matter = extract_front_matter(text)
        if front_matter:
            meta_info["files"][file_path.name] = front_matter
        cleaned = remove_front_matter(text).strip()
        combined_parts.append(cleaned)
        meta_info["order"].append(file_path.name)

    Path(output_combined).write_text("\n\n".join(combined_parts) + "\n", encoding='utf-8')
    Path(output_meta_json).write_text(json.dumps(meta_info, indent=2), encoding='utf-8')

    print(f"Combined file saved as: {output_combined}")
    print(f"Metadata JSON saved as: {output_meta_json}")

def write_with_safety(path, content, force=False, backup=False):
    """Write a file safely with force/backup options and diff-aware prompt."""
    if path.exists():
        current = path.read_text(encoding='utf-8')
        if current == content:
            print(f"No changes for: {path}")
            return

        if backup and not force:
            backup_path = path.with_suffix(path.suffix + ".bak")
            shutil.copy2(path, backup_path)
            print(f"Backup created: {backup_path}")
        elif not force:
            while True:
                choice = input(f"{path} has changed. [y] overwrite, [b] backup+overwrite, [n] skip? ").strip().lower()
                if choice in ("y", "yes"):
                    break
                elif choice in ("b", "backup"):
                    backup_path = path.with_suffix(path.suffix + ".bak")
                    shutil.copy2(path, backup_path)
                    print(f"Backup created: {backup_path}")
                    break
                elif choice in ("n", "no", ""):
                    print(f"Skipped: {path}")
                    return
                else:
                    print("Please choose [y], [b], or [n].")
    path.write_text(content, encoding='utf-8')
    print(f"Wrote: {path}")

def split_with_front_matter(input_combined, output_dir, metadata_file, force=False, backup=False):
    """Split a combined markdown file back into original files, restoring front matter."""
    combined_text = Path(input_combined).read_text(encoding='utf-8')

    if not Path(metadata_file).exists():
        print(f"Metadata file not found: {metadata_file}")
        sys.exit(1)

    meta_info = json.loads(Path(metadata_file).read_text(encoding='utf-8'))
    os.makedirs(output_dir, exist_ok=True)

    order = meta_info.get("order", [])
    frontmatters = meta_info.get("files", {})

    # Split by H1 or H2 headings — each original file should start with an H1 or H2 heading.
    sections = re.split(r'(?=^(?:#{1,2})\s)', combined_text, flags=re.MULTILINE)
    sections = [s.strip() for s in sections if s.strip()]

    if len(sections) != len(order):
        raise Exception(f"{len(sections)} sections found but {len(order)} files listed.")
    else:
        chunks = sections

    for i, filename in enumerate(order):
        output_path = Path(output_dir, filename)
        body = chunks[i].strip() if i < len(chunks) else ""

        # Restore front matter if available
        front_matter = frontmatters.get(filename)
        if front_matter:
            content = f"---\n{front_matter}\n---\n\n{body}\n"
        else:
            content = body + "\n"

        write_with_safety(output_path, content, force=force, backup=backup)

    print(f"Split complete. {len(order)} files processed.")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage:")
        print("  Combine: python manage_docs.py combine")
        print("  Split:   python manage_docs.py split [--force|--backup]")
        sys.exit(1)

    command = sys.argv[1].lower()
    force = "--force" in sys.argv
    backup = "--backup" in sys.argv

    if command == "combine":
        combine_markdown(filenames, "stack_lang_spec.md", "metadata.json")

    elif command == "split":
        split_with_front_matter("stack_lang_spec.md", "docs", "metadata.json",
                                force=force, backup=backup)

    else:
        print("Unknown command. Use 'combine' or 'split'.")