import sys
import json
import yaml
import argparse
import os
from copy import deepcopy
from importlib.metadata import version, PackageNotFoundError

try:
    __version__ = version("adp-api-transformer")
except PackageNotFoundError:
    __version__ = "unknown"

def load_spec(file_path):
    with open(file_path, 'r') as f:
        if file_path.endswith(('.yaml', '.yml')):
            return yaml.safe_load(f)
        elif file_path.endswith('.json'):
            return json.load(f)
        else:
            raise ValueError("Unsupported file format. Please provide a .json or .yaml file.")

def generate_operation_id(method, path):
    clean_path = path.strip('/').replace('/', '_').replace('{', '').replace('}', '')
    return f"{method}_{clean_path}"+"-{{consumerphase}}"

def add_missing_attributes(spec):
    if 'info' not in spec or not isinstance(spec['info'], dict):
        spec['info'] = {}
    spec['info']['title'] = '{{apiTitle}}'

    paths = spec.get('paths', {})
    for path, methods in paths.items():
        for method, operation in methods.items():
            if not isinstance(operation, dict):
                continue
            if 'operationId' not in operation or not operation['operationId']:
                operation['operationId'] = generate_operation_id(method, path)
            elif not operation['operationId'].endswith('-{{consumerphase}}'):
                operation['operationId'] = operation['operationId'] + '-{{consumerphase}}'
            if 'summary' not in operation:
                operation['summary'] = f"{method.upper()} {path}"
            # Add a default tag if missing, using the first path segment
            if 'tags' not in operation or not operation['tags']:
                first_segment = path.strip('/').split('/')[0]
                operation['tags'] = [first_segment]
    return spec

def save_as_json(spec, output_path):
    with open(output_path, 'w') as f:
        json.dump(spec, f, indent=2)

def get_referenced_schemas(spec_part):
    """
    Recursively finds all schema references in a spec part.
    Returns a set of schema names referenced via $ref.
    """
    refs = set()
    
    def extract_refs(obj):
        if isinstance(obj, dict):
            if '$ref' in obj:
                ref = obj['$ref']
                if '#/components/schemas/' in ref:
                    schema_name = ref.split('#/components/schemas/')[-1]
                    refs.add(schema_name)
            for value in obj.values():
                extract_refs(value)
        elif isinstance(obj, list):
            for item in obj:
                extract_refs(item)
    
    extract_refs(spec_part)
    return refs

def get_all_referenced_schemas(schemas, initial_refs):
    """
    Recursively resolves all referenced schemas.
    If schema A references schema B, both are included.
    """
    all_refs = set(initial_refs)
    to_process = list(initial_refs)
    
    while to_process:
        schema_name = to_process.pop(0)
        if schema_name in schemas:
            schema = schemas[schema_name]
            refs = get_referenced_schemas(schema)
            for ref in refs:
                if ref not in all_refs:
                    all_refs.add(ref)
                    to_process.append(ref)
    
    return all_refs

def split_openapi_by_tag(spec):
    """
    Splits the OpenAPI spec into multiple specs based on the first tag of each operation.
    The path in each split spec starts after the tag segment, or '/' if the tag is the last segment.
    Only includes relevant schemas for each tag, plus general/shared schemas.
    """
    tag_specs = {}
    all_schemas = spec.get("components", {}).get("schemas", {})
    
    # Define general/shared schema names that should be included in all specs
    general_schemas = {"security", "error", "paginated", "metadata"}

    for path, methods in spec.get("paths", {}).items():
        for method, operation in methods.items():
            if not isinstance(operation, dict):
                continue
            tags = operation.get("tags", [])
            if not tags:
                continue
            tag = tags[0]
            tag_lower = tag.lower()
            parts = path.strip('/').split('/')
            try:
                tag_index = [p.lower() for p in parts].index(tag_lower)
                # If tag is the last segment, use '/'
                if tag_index == len(parts) - 1:
                    new_path = '/'
                else:
                    # new_path is everything after the tag (exclusive)
                    new_parts = parts[tag_index + 1 :]
                    new_path = '/' + '/'.join(new_parts)
            except ValueError:
                new_path = path

            if tag not in tag_specs:
                components = spec.get('components', {})
                # Create a minimal tag spec with only top-level structure
                tag_specs[tag] = {
                    "openapi": spec.get("openapi"),
                    "info": deepcopy(spec.get("info", {})),
                    "paths": {},
                    "components": {
                        "schemas": {},
                        "securitySchemes": deepcopy(components.get('securitySchemes', {}))
                    }
                }
                # Copy other top-level fields if they exist
                for key in spec:
                    if key not in ["openapi", "info", "paths", "components"]:
                        tag_specs[tag][key] = deepcopy(spec[key])

            if new_path not in tag_specs[tag]["paths"]:
                tag_specs[tag]["paths"][new_path] = {}
            tag_specs[tag]["paths"][new_path][method] = deepcopy(operation)

    # Filter schemas for each tag
    for tag, tag_spec in tag_specs.items():
        # Find all referenced schemas in this tag's paths
        referenced_refs = set()
        for path_obj in tag_spec.get("paths", {}).values():
            for operation in path_obj.values():
                if isinstance(operation, dict):
                    refs = get_referenced_schemas(operation)
                    referenced_refs.update(refs)
        
        # Add general schemas to the referenced set
        referenced_refs.update(general_schemas)
        
        # Get all schemas including nested references
        all_referenced = get_all_referenced_schemas(all_schemas, referenced_refs)
        
        # Only include referenced schemas
        tag_spec["components"]["schemas"] = {
            name: deepcopy(all_schemas[name]) for name in all_referenced if name in all_schemas
        }

    return tag_specs

def split_openapi_by_path(spec, ignore_prefix=''):
    """
    Splits the OpenAPI spec into multiple specs based on the first path segment
    after stripping an optional prefix.
    E.g. with ignore_prefix='api/v1', '/api/v1/studies/{id}' groups under 'studies'
    and the resulting path becomes '/{id}'.
    """
    group_specs = {}
    all_schemas = spec.get("components", {}).get("schemas", {})
    general_schemas = {"security", "error", "paginated", "metadata"}
    prefix_parts = [p for p in ignore_prefix.strip('/').split('/') if p]

    for path, methods in spec.get("paths", {}).items():
        parts = [p for p in path.strip('/').split('/') if p]

        # Strip the ignore_prefix from the beginning
        remaining = parts
        if prefix_parts:
            if [p.lower() for p in parts[:len(prefix_parts)]] == [p.lower() for p in prefix_parts]:
                remaining = parts[len(prefix_parts):]

        if not remaining:
            group = 'root'
            new_path = '/'
        else:
            group = remaining[0]
            if len(remaining) == 1:
                new_path = '/'
            else:
                new_path = '/' + '/'.join(remaining[1:])

        for method, operation in methods.items():
            if not isinstance(operation, dict):
                continue

            if group not in group_specs:
                components = spec.get('components', {})
                group_specs[group] = {
                    "openapi": spec.get("openapi"),
                    "info": deepcopy(spec.get("info", {})),
                    "paths": {},
                    "components": {
                        "schemas": {},
                        "securitySchemes": deepcopy(components.get('securitySchemes', {}))
                    }
                }
                for key in spec:
                    if key not in ["openapi", "info", "paths", "components"]:
                        group_specs[group][key] = deepcopy(spec[key])

            if new_path not in group_specs[group]["paths"]:
                group_specs[group]["paths"][new_path] = {}
            group_specs[group]["paths"][new_path][method] = deepcopy(operation)

    # Filter schemas for each group
    for group, group_spec in group_specs.items():
        referenced_refs = set()
        for path_obj in group_spec.get("paths", {}).values():
            for operation in path_obj.values():
                if isinstance(operation, dict):
                    refs = get_referenced_schemas(operation)
                    referenced_refs.update(refs)
        referenced_refs.update(general_schemas)
        all_referenced = get_all_referenced_schemas(all_schemas, referenced_refs)
        group_spec["components"]["schemas"] = {
            name: deepcopy(all_schemas[name]) for name in all_referenced if name in all_schemas
        }

    return group_specs

def save_split_specs(spec, output_dir):
    """
    Splits the OpenAPI spec and saves each tag spec as a JSON file in the output directory.
    """
    os.makedirs(output_dir, exist_ok=True)
    tag_specs = split_openapi_by_tag(spec)
    for tag, tag_spec in tag_specs.items():
        filename = f"{tag.lower()}.openapi.json"
        filepath = os.path.join(output_dir, filename)
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(tag_spec, f, indent=2)

def save_split_by_path_specs(spec, output_dir, ignore_prefix=''):
    """
    Splits the OpenAPI spec by path and saves each group as a JSON file.
    """
    os.makedirs(output_dir, exist_ok=True)
    group_specs = split_openapi_by_path(spec, ignore_prefix)
    for group, group_spec in group_specs.items():
        filename = f"{group.lower()}.openapi.json"
        filepath = os.path.join(output_dir, filename)
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(group_spec, f, indent=2)

def convert_api_spec():
    parser = argparse.ArgumentParser(description="Modify OpenAPI spec and output as JSON, or split by tag.")
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    parser.add_argument("input_file", help="Path to the input OpenAPI spec file (JSON or YAML).")
    parser.add_argument("output", nargs='?', default="openapi.json", help="Output file or directory.")
    parser.add_argument("--split-by-tag", action="store_true", help="Split the OpenAPI spec by tag into multiple files.")
    parser.add_argument("--split-by-path", action="store_true", help="Split the OpenAPI spec by the first path segment into multiple files.")
    parser.add_argument("--ignore-prefix", default="", help="Path prefix to ignore when splitting by path (e.g. 'api/v1').")
    args = parser.parse_args()
    input_path = os.path.abspath(args.input_file)
    output_path = os.path.abspath(args.output)
    spec = load_spec(input_path)
    modified_spec = add_missing_attributes(spec)
    output_dir = os.path.splitext(input_path)[0]
    if args.split_by_tag:
        save_split_specs(modified_spec, output_dir)
        print(f"OpenAPI spec split by tag and saved to directory {output_dir}")
    elif args.split_by_path:
        save_split_by_path_specs(modified_spec, output_dir, args.ignore_prefix)
        print(f"OpenAPI spec split by path and saved to directory {output_dir}")
    else:
        save_as_json(modified_spec, output_path)
        print(f"Modified OpenAPI spec saved to {output_path}")