"""
Main SDK generation orchestration.
"""

import os
from typing import Dict, List, Any, Set, Union, cast, Optional
import traceback
import yaml

from .schema import (
    get_response_type,
    identify_response_schemas,
    build_dependency_graph,
    get_request_body_parameters,
    topological_sort,
)
from .types import to_snake_case, get_type_from_schema
from .template import render_template


def generate_sdk(
    openapi_spec: Union[str, Dict[str, Any]],
    output_dir: str,
    package_name: Optional[str] = None,
    py_version: float = 3.10,
    use_modern_py: bool = False,
    template_dir: Optional[str] = None,
) -> bool:
    """
    Generate an SDK from an OpenAPI specification.

    Args:
        openapi_spec: Path to an OpenAPI spec file or dict containing the spec
        output_dir: Directory where the SDK will be generated
        package_name: Name of the package (default: derived from spec title)
        py_version: Python version to target (default: 3.10)
        use_modern_py: Use modern Python features (default: False)
        template_dir: Directory containing custom templates (default: None)

    Returns:
        bool: True if successful, False otherwise
    """
    try:
        # If openapi_spec is a string (file path), load it
        if isinstance(openapi_spec, str):
            openapi_spec = load_openapi_spec(openapi_spec)

        # Handle case where spec loading fails
        if openapi_spec is None:
            print("Failed to load OpenAPI spec")
            return False

        # Create output directory
        os.makedirs(output_dir, exist_ok=True)

        # Get package name from spec if not provided
        if not package_name:
            info = openapi_spec.get("info", {})
            title = info.get("title", "api")
            package_name = title.lower().replace(" ", "_")

        # Extract schemas (components/schemas)
        schemas_dict = get_schemas_from_spec(openapi_spec)

        # Get schema names
        list(schemas_dict.keys())

        # Build dependency graph - Pass only schemas_dict
        dependency_graph = build_dependency_graph(schemas_dict)

        # Get topologically sorted schema names
        sorted_schemas = topological_sort(dependency_graph, schemas_dict)

        # Identify response schemas (returns a Set)
        response_schemas_set = identify_response_schemas(openapi_spec)

        # Generate response models - Pass the set
        assert package_name is not None
        assert isinstance(package_name, str)
        response_models_info = generate_response_models(
            openapi_spec,
            schemas_dict,
            sorted_schemas,
            output_dir,
            package_name,
            use_modern_py,
            response_schemas_set,  # Pass the set
        )

        # Generate client
        client_data = prepare_client_data(
            openapi_spec, output_dir, package_name, use_modern_py, response_models_info
        )  # Pass the dict returned by generate_response_models

        # Render client template
        client_code = render_template("client.py.j2", client_data, template_dir)

        # Create package directory if needed
        package_dir = os.path.join(output_dir, package_name)
        os.makedirs(package_dir, exist_ok=True)

        # Write client file
        with open(os.path.join(package_dir, "client.py"), "w") as f:
            f.write(client_code)

        # Write __init__.py for package
        with open(os.path.join(package_dir, "__init__.py"), "w") as f:
            f.write(f'"""SDK for {package_name} API.\n\nGenerated by Damascus."""\n\n')
            f.write("from .client import *\n")

        # Write top-level __init__.py
        with open(os.path.join(output_dir, "__init__.py"), "w") as f:
            f.write(f'"""SDK for {package_name} API.\n\nGenerated by Damascus."""\n\n')
            f.write(f"from {package_name} import *\n")

        return True
    except Exception as e:
        print(f"Error generating SDK: {e}")
        traceback.print_exc()
        return False


def load_openapi_spec(openapi_spec_path: str, headers: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
    """
    Load an OpenAPI specification from a file path or URL.

    Args:
        openapi_spec_path: Path to an OpenAPI file or URL
        headers: Optional headers to use when fetching from a URL

    Returns:
        The loaded OpenAPI specification as a dictionary

    Raises:
        ValueError: If the spec couldn't be loaded or parsed
    """
    # Check if the spec path is a URL
    if openapi_spec_path.startswith(("http://", "https://")):
        import requests

        try:
            response = requests.get(openapi_spec_path, headers=headers)
            response.raise_for_status()
            if openapi_spec_path.endswith(".yaml") or openapi_spec_path.endswith(".yml"):
                return cast(Dict[str, Any], yaml.safe_load(response.text))
            else:
                return cast(Dict[str, Any], response.json())
        except Exception as e:
            raise ValueError(f"Error loading OpenAPI spec from URL: {e}")
    else:
        # Load from file
        try:
            with open(openapi_spec_path, "r") as f:
                if openapi_spec_path.endswith(".yaml") or openapi_spec_path.endswith(".yml"):
                    # Cast the result of yaml.safe_load
                    return cast(Dict[str, Any], yaml.safe_load(f))
                else:
                    import json

                    # Cast the result of json.load
                    return cast(Dict[str, Any], json.load(f))
        except Exception as e:
            raise ValueError(f"Error loading OpenAPI spec from file: {e}")


def get_schemas_from_spec(openapi_spec: Dict[str, Any]) -> Dict[str, Any]:
    """
    Extract schema definitions from an OpenAPI specification.

    Args:
        openapi_spec: The OpenAPI specification as a dictionary

    Returns:
        A dictionary mapping schema names to their definitions
    """
    if "components" in openapi_spec and "schemas" in openapi_spec["components"]:
        return cast(Dict[str, Any], openapi_spec["components"]["schemas"])
    return {}


def generate_response_models(
    openapi_spec: Dict[str, Any],
    schemas_dict: Dict[str, Any],
    sorted_schemas: List[str],
    output_dir: str,
    package_name: str,
    use_modern_py: bool = False,
    response_schemas_set: Optional[Set[str]] = None,  # Expect a Set
) -> Dict[str, Dict[str, Any]]:  # Keep returning Dict for now, might change later
    """
    Generate model classes for response schemas.

    Args:
        openapi_spec: The OpenAPI specification
        schemas_dict: Dictionary of schemas from the OpenAPI spec
        sorted_schemas: List of schema names in topological order
        output_dir: Directory to output the generated code
        package_name: Name of the package
        use_modern_py: Whether to use modern Python features
        response_schemas_set: Set of response schema names

    Returns:
        Dict[str, Dict[str, Any]]: Dictionary mapping schema names to response schema information
    """
    # Create models directory
    package_dir = os.path.join(output_dir, package_name)
    models_dir = os.path.join(package_dir, "models")
    os.makedirs(models_dir, exist_ok=True)

    # Create __init__.py in models directory
    with open(os.path.join(models_dir, "__init__.py"), "w") as f:
        f.write('"""Response model classes.\n\nAutomatically generated by Damascus."""\n\n')

    # Identify response schemas - No longer needed here, passed as arg
    # response_schemas = identify_response_schemas(openapi_spec)
    response_schemas_set = response_schemas_set or set()  # Ensure it's a set

    # Only process schemas that are used in responses (check against the set)
    schema_names_to_process = [name for name in sorted_schemas if name in response_schemas_set]

    # Generate model code for each schema
    processed_schemas = {}
    for schema_name in schema_names_to_process:
        schema = schemas_dict.get(schema_name)
        if schema:
            model_code = generate_model_code(schema_name, schema, schemas_dict)

            # Write model file
            with open(os.path.join(models_dir, f"{to_snake_case(schema_name)}.py"), "w") as f:
                f.write(model_code)

            # Add import to __init__.py
            with open(os.path.join(models_dir, "__init__.py"), "a") as f:
                f.write(f"from .{to_snake_case(schema_name)} import {schema_name}\n")

            processed_schemas[schema_name] = schema  # Store processed schema info

    return processed_schemas  # Return the dict of processed schemas


def generate_model_code(schema_name: str, schema: Dict[str, Any], schemas: Dict[str, Any]) -> str:
    """
    Generate Python dataclass code for a model.

    Args:
        schema_name: Name of the model
        schema: Schema of the model
        schemas: All component schemas

    Returns:
        Python code for the model as a string
    """
    properties = schema.get("properties", {})
    required_props = schema.get("required", [])

    required_fields = []
    optional_fields = []

    for prop_name, prop_schema in properties.items():
        field_name = to_snake_case(prop_name)
        field_type = get_type_from_schema(prop_schema, schemas)
        field_required = prop_name in required_props

        # Prepare field definition
        field_def = {"name": field_name, "type": field_type, "required": field_required}

        if field_required:
            required_fields.append(field_def)
        else:
            optional_fields.append(field_def)

    # Sort required fields before optional fields for Python dataclass
    fields = required_fields + optional_fields

    # Import statements for types used in the model
    imports = ["from dataclasses import dataclass, field"]
    # Base imports needed for Optional fields and potentially complex types
    typing_imports = {"Optional", "Dict", "Any"}

    for field_info in fields:
        # Cast the type to string, as get_type_from_schema should return a string representation
        field_type = str(field_info["type"])
        # Ensure field_type is a string before checking for 'List' substring
        if "List" in field_type:
            typing_imports.add("List")
        # Consider adding checks for other complex types like Union if needed in the future

    # Construct the import string from the set of required types
    if typing_imports:
        import_str = ", ".join(sorted(list(typing_imports)))
        imports.append(f"from typing import {import_str}")

    class_name = schema_name

    # Build the model code
    code = "\n".join(imports) + "\n\n\n"
    code += f"@dataclass\nclass {class_name}:\n"

    if not fields:
        code += "    pass\n"
    else:
        for field in fields:
            # Explicitly cast dictionary lookups to expected types for mypy
            field_name = str(field["name"])
            field_type = str(field["type"])
            field_required = bool(field["required"])

            if field_required:
                code += f"    {field_name}: {field_type}\n"
            else:
                code += f"    {field_name}: Optional[{field_type}] = None\n"

    return code


def create_models_init(model_mapping: Dict[str, str], models_dir: str) -> None:
    """
    Create __init__.py file for the models package.

    Args:
        model_mapping: Mapping of schema names to class names
        models_dir: Directory path for models
    """
    exports = []

    for schema_name, class_name in model_mapping.items():
        exports.append(class_name)

    code = []
    code.append('"""')
    code.append("Response model classes.")
    code.append('"""\n')
    code.append("from .models import *")
    code.append("\n")
    code.append(f"__all__ = {repr(exports)}")

    # Write the file
    init_path = os.path.join(models_dir, "__init__.py")
    with open(init_path, "w") as f:
        f.write("\n".join(code))


def prepare_client_data(
    openapi_spec: Dict[str, Any],
    output_dir: str,
    package_name: str,
    use_modern_py: bool = False,
    response_schemas: Optional[Dict[str, Dict[str, Any]]] = None,
) -> Dict[str, Any]:
    """
    Prepare data structure for the client template.

    Args:
        openapi_spec: The OpenAPI specification
        output_dir: Directory to output the generated code
        package_name: Name of the package
        use_modern_py: Whether to use modern Python features
        response_schemas: Dictionary of response schemas

    Returns:
        Dict with data needed for rendering the client template
    """
    response_schemas = response_schemas or {}

    # Get basic info from the spec
    info = openapi_spec.get("info", {})
    title = info.get("title", "API")
    version = info.get("version", "1.0.0")
    description = info.get("description", f"SDK for {title} API")

    # Extract server URL
    servers = openapi_spec.get("servers", [{"url": "https://api.example.com"}])
    server_url = servers[0]["url"] if servers else "https://api.example.com"

    # Process security schemes
    security_schemes = {}
    if "components" in openapi_spec and "securitySchemes" in openapi_spec["components"]:
        security_schemes = openapi_spec["components"]["securitySchemes"]

    # Process paths
    paths: Dict[str, Dict[str, Dict[str, Any]]] = {}
    for path, path_item in openapi_spec.get("paths", {}).items():
        for method, operation in path_item.items():
            if method not in ["get", "post", "put", "delete", "patch"]:
                continue

            operation_id = operation.get("operationId")
            if not operation_id:
                # Generate operation ID from path and method
                path_parts = [p for p in path.split("/") if p]
                operation_id = method + "_" + "_".join(path_parts)
                operation_id = to_snake_case(operation_id)

            # Process parameters
            parameters = []
            for param in operation.get("parameters", []):
                param_schema = param.get("schema", {})
                param_type = get_type_from_schema(param_schema, openapi_spec.get("components", {}).get("schemas", {}))

                parameters.append(
                    {
                        "name": to_snake_case(param["name"]),
                        "required": param.get("required", False),
                        "in": param.get("in"),
                        "type": param_type,
                        "description": param.get("description", ""),
                    }
                )

            # Process request body
            request_body_params = []
            if "requestBody" in operation:
                request_body = operation["requestBody"]
                request_body_params = get_request_body_parameters(request_body, openapi_spec.get("components", {}).get("schemas", {}))

            # Process response
            response_type = "dict"  # Default response type
            response_model = None

            for status_code, response in operation.get("responses", {}).items():
                if int(status_code) < 400:  # Only care about successful responses
                    response_schema, is_array = get_response_type(response, openapi_spec.get("components", {}).get("schemas", {}))
                    if response_schema:
                        # Extract type from the schema using get_type_from_schema
                        response_type_str = get_type_from_schema(
                            response_schema,
                            openapi_spec.get("components", {}).get("schemas", {}),
                        )
                        if is_array:
                            response_type = f"List[{response_type_str}]"
                        else:
                            response_type = response_type_str

                        # Check if the type is a model reference
                        if "$ref" in response_schema:
                            response_model = response_schema["$ref"].split("/")[-1]
                        elif is_array and "items" in response_schema and "$ref" in response_schema["items"]:
                            response_model = response_schema["items"]["$ref"].split("/")[-1]
                        else:
                            response_model = None  # Not a direct model reference

                        break

            # Store operation details
            path_key = path
            if path_key not in paths:
                paths[path_key] = {}

            paths[path_key][method] = {
                "operation_id": operation_id,
                "parameters": parameters,
                "request_body_params": request_body_params,
                "response_type": response_type,
                "response_model": response_model,
                "summary": operation.get("summary", ""),
                "description": operation.get("description", ""),
            }

    return {
        "title": title,
        "version": version,
        "description": description,
        "server_url": server_url,
        "paths": paths,
        "security_schemes": security_schemes,
        "use_modern_py": use_modern_py,
        "package_name": package_name,
        "has_response_models": bool(response_schemas),  # Check the dict passed in
    }


def generate_client_file(output_dir: str, client_data: Dict[str, Any]) -> None:
    """
    Generates the client file from the given client data.

    Args:
        output_dir: Directory where the SDK code will be generated
        client_data: Data to pass to the client template
    """
    # Render the client template
    client_code = render_template("client.py.j2", client_data)

    # Write the generated code to file
    client_file_path = os.path.join(output_dir, "__init__.py")  # client in __init__.py
    with open(client_file_path, "w") as f:
        f.write(client_code)


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        # If a file path is provided as an argument
        openapi_file = sys.argv[1]
        print(f"Generating SDK from: {openapi_file}")
        # Add output_dir argument
        generate_sdk(load_openapi_spec(openapi_file), output_dir="./generated_sdk")
    else:
        # Default behavior
        print("No OpenAPI file specified, using default: openapi.json")
        print("Usage: python -m damascus.core [path_to_openapi_json or URL]")
        # Add output_dir argument
        generate_sdk(load_openapi_spec("openapi.json"), output_dir="./generated_sdk")
