import json
import os
import re
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Dict, List, Optional

import datamodel_code_generator
from datamodel_code_generator import DataModelType, InputFileType, generate
from pydantic import BaseModel

from oold.utils.codegen import OOLDJsonSchemaParserFixedRefs as OOLDJsonSchemaParser


class Generator:
    class GenerateParams(BaseModel):
        json_schemas: List[Dict]
        """JSON SCHEMA source(s)"""
        preprocess: bool = True
        """Preprocess the JSON schemas before generating the models"""
        main_schema: Optional[str] = None
        """File name of the main schema"""
        output_model_type: Optional[DataModelType] = (
            DataModelType.PydanticV2BaseModel,
        )
        """Output model type, e.g. PydanticV2BaseModel or PydanticBaseModel"""
        output_model_path: Optional[Path] = (
            Path(__file__).parent / "model" / "example.py"
        )
        """Output model path, if not set the model will be generated
        in the current directory"""
        working_dir_path: Optional[Path] = None
        """Working directory to store intermedia files
        and the generated partial models"""
        generate_init_py_files: bool = True
        """Generate __init__.py files along the output_model_path"""

    def generate(
        self,
        params: GenerateParams,
    ):
        if params.preprocess:
            self.preprocess(
                Generator.PreprocessParams(json_schemas=params.json_schemas)
            )

        # monkey patch class
        datamodel_code_generator.parser.jsonschema.JsonSchemaParser = (
            OOLDJsonSchemaParser
        )

        with TemporaryDirectory() as temporary_directory_name:
            temporary_directory = Path(temporary_directory_name)
            if params.working_dir_path is not None:
                temporary_directory = params.working_dir_path

            input = Path(temporary_directory)
            if params.main_schema is not None:
                input = Path(temporary_directory / Path(params.main_schema))

            output = params.output_model_path
            if params.generate_init_py_files:
                # generate __init__.py files in every subdirectory
                # of the output model path that does not exist yet
                # output may be a file or a directory

                # check if output is a file or a directory
                target_dir = output
                if params.main_schema is not None:
                    target_dir = output.parent

                # interate over the target_dir path, starting at the top level dir
                # e.g. 'C:' or '/var'
                for segment in target_dir.parts:
                    # create the segment path
                    segment_path = Path(
                        *target_dir.parts[: target_dir.parts.index(segment) + 1]
                    )
                    # check if the segment path exists
                    if not segment_path.exists():
                        # create the __init__.py file
                        os.makedirs(segment_path, exist_ok=False)
                        init_file = segment_path / "__init__.py"
                        with open(init_file, "w", encoding="utf-8") as f:
                            f.write("# Generated by oold.generator\n")

            for schema in params.json_schemas:
                name = schema["id"]
                os.makedirs(
                    os.path.dirname(Path(temporary_directory / (name + ".json"))),
                    exist_ok=True,
                )
                with open(
                    Path(temporary_directory / (name + ".json")), "w", encoding="utf-8"
                ) as f:
                    schema_str = json.dumps(
                        schema, ensure_ascii=False, indent=4
                    ).replace("dollarref", "$ref")
                    # print(schema_str)
                    f.write(schema_str)

            if params.output_model_type == DataModelType.PydanticV2BaseModel:
                base_class = "oold.model.LinkedBaseModel"
            else:
                base_class = "oold.model.v1.LinkedBaseModel"
            generate(
                input_=input,
                # json_schema,
                input_file_type=InputFileType.JsonSchema,
                # input_filename="Foo.json",
                output=output,
                # set up the output model types
                output_model_type=params.output_model_type,
                # custom_template_dir=Path(model_dir_path),
                field_include_all_keys=True,
                base_class=base_class,
                # use_default = True,
                enum_field_as_literal="all",
                use_title_as_name=True,
                use_schema_description=True,
                use_field_description=True,
                encoding="utf-8",
                use_double_quotes=True,
                collapse_root_models=True,
                reuse_model=True,
                # create MyEnum(str, Enum) instead of MyEnum(Enum)
                use_subclass_enum=True,
                additional_imports=["pydantic.ConfigDict"]
                if params.output_model_type == DataModelType.PydanticV2BaseModel
                else [],
                apply_default_values_for_required_fields=True,
            )

            if params.main_schema is not None:
                content = ""
                with open(output, "r", encoding="utf-8") as f:
                    content = f.read()
                os.remove(output)

                content = re.sub(
                    r"(UUID = Field\(...)",
                    r"UUID = Field(default_factory=uuid4",
                    content,
                )  # enable default value for uuid

                if params.output_model_type == DataModelType.PydanticBaseModel:
                    # we are now using pydantic.v1
                    # pydantic imports lead to uninitialized fields
                    # (FieldInfo still present)
                    content = re.sub(
                        r"(from pydantic import)", "from pydantic.v1 import", content
                    )

                # write the content to the file
                with open(output, "w", encoding="utf-8") as f:
                    f.write(content)

    class PreprocessParams(BaseModel):
        json_schemas: List[Dict]
        """JSON SCHEMA source(s)"""

    def preprocess(self, params: PreprocessParams):
        for schema in params.json_schemas:
            # schema = self.merge_property_schemas(schema)
            for property_key in schema.get("properties", {}):
                property = schema["properties"][property_key]
                if "range" in property:
                    if "type" in property:
                        del property["type"]
                    # if range is a string we create a allOf with a ref to the range
                    if isinstance(property["range"], str):
                        property["allOf"] = [{"$ref": property["range"]}]
                    else:
                        property["$ref"] = property["range"]
                    if "required" in schema and property_key in schema["required"]:
                        # if no default value is set, remove the property from required
                        if "default" not in property:
                            schema["required"].remove(property_key)
                        if "x-oold-required-iri" not in property:
                            property["x-oold-required-iri"] = True
                if "items" in property:
                    if "range" in property["items"]:
                        if "type" in property["items"]:
                            del property["items"]["type"]
                        if isinstance(property["items"]["range"], str):
                            property["items"]["allOf"] = [
                                {"$ref": property["items"]["range"]}
                            ]
                        else:
                            property["items"]["$ref"] = property["items"]["range"]
                        property["range"] = property["items"]["range"]
                        if "required" in schema and property_key in schema["required"]:
                            # if no default value is set,
                            # remove the property from required
                            if "default" not in property["items"]:
                                schema["required"].remove(property_key)
                            if "x-oold-required-iri" not in property:
                                property["x-oold-required-iri"] = True

                    if "properties" in property["items"]:
                        self.preprocess(
                            Generator.PreprocessParams(json_schemas=[property["items"]])
                        )

                if "properties" in property:
                    self.preprocess(Generator.PreprocessParams(json_schemas=[property]))
