import os
import re
from rdflib import Graph, Namespace, RDF, RDFS, URIRef
from typing import Dict, List, Set, Optional

CIM = Namespace("http://iec.ch/TC57/2013/CIM-schema-cim16#")
CIMS = Namespace("http://iec.ch/TC57/1999/rdf-schema-extensions-19990926#")

class SchemaParser:
    def __init__(self, rdfs_dir: str):
        self.rdfs_dir = rdfs_dir
        self.all_classes = {} # name -> data
        self.profiles = {} # profile_name -> {classes: [], properties: []}
        
    def parse(self):
        # First pass: collect all classes and their "primary" profile
        # We'll use the file where the RDFS.Class is defined as the primary.
        for filename in sorted(os.listdir(self.rdfs_dir)):
            if not filename.endswith(".rdf"): continue
            
            g = Graph()
            g.parse(os.path.join(self.rdfs_dir, filename), format="xml")
            
            profile_name = self._get_profile_name(filename)
            if profile_name not in self.profiles:
                self.profiles[profile_name] = {"classes": [], "properties": []}
            
            for s in g.subjects(RDF.type, RDFS.Class):
                name = str(s).split("#")[-1]
                if name not in self.all_classes:
                    parent = g.value(s, RDFS.subClassOf)
                    parent_name = str(parent).split("#")[-1] if parent else None
                    stereotype = g.value(s, CIMS.stereotype)
                    
                    self.all_classes[name] = {
                        "uri": s,
                        "parent": parent_name,
                        "is_enum": str(stereotype) == "enumeration",
                        "properties": {}, # uri -> data
                        "profile": profile_name,
                        "enum_values": set()
                    }
                self.profiles[profile_name]["classes"].append(name)
                
                if self.all_classes[name]["is_enum"]:
                    for val_s in g.subjects(RDF.type, s):
                        val_name = str(val_s).split(".")[-1]
                        self.all_classes[name]["enum_values"].add(val_name)

            for s in g.subjects(RDF.type, RDF.Property):
                name = str(s).split("#")[-1]
                domain = g.value(s, RDFS.domain)
                if not domain: continue
                
                domain_name = str(domain).split("#")[-1]
                range_val = g.value(s, RDFS.range)
                data_type = g.value(s, CIMS.dataType)
                prop_range = range_val or data_type
                range_name = str(prop_range).split("#")[-1] if prop_range else "Any"
                multiplicity = g.value(s, CIMS.multiplicity)
                mult_str = str(multiplicity).split("#M:")[-1] if multiplicity else "0..1"
                
                prop_data = {
                    "name": name.split(".")[-1],
                    "uri": str(s),
                    "range": range_name,
                    "is_list": ".." in mult_str and not mult_str.endswith("1"),
                    "domain": domain_name
                }
                
                # Add property to the class (even if class was defined in another profile)
                if domain_name in self.all_classes:
                    self.all_classes[domain_name]["properties"][str(s)] = prop_data

    def _get_profile_name(self, filename):
        name = re.sub(r'ProfileRDFSAugmented.*', '', filename)
        name = re.sub(r'RDFSAugmented.*', '', name)
        return name.replace('-', '_').strip('_')

    def generate(self, out_dir: str):
        models_dir = os.path.join(out_dir, "models")
        os.makedirs(models_dir, exist_ok=True)
        
        type_map = {
            "String": "str", "Integer": "int", "Float": "float", "Boolean": "bool",
            "DateTime": "str", "Date": "str", "Decimal": "float", "Seconds": "float",
            "Voltage": "float", "ActivePower": "float", "ReactivePower": "float",
            "ApparentPower": "float", "CurrentFlow": "float", "Resistance": "float",
            "Reactance": "float", "Susceptance": "float", "Conductance": "float",
            "PerCent": "float", "AngleDegrees": "float", "Length": "float", "Any": "Any"
        }

        module_header = [
            "from __future__ import annotations",
            "from typing import List, Optional, Union, Any",
            "from enum import Enum",
            "from pydantic import Field",
            "from ..base import CIMBase, CIMModelRegistry",
            "from rdflib import URIRef",
            "",
            "CIM = URIRef('http://iec.ch/TC57/2013/CIM-schema-cim16#')",
            ""
        ]

        def class_block(name: str):
            data = self.all_classes[name]
            lines = []
            if data["is_enum"]:
                lines.append(f"class {name}(str, Enum):")
                if not data["enum_values"]:
                    lines.append("    UNKNOWN = 'UNKNOWN'")
                for val in sorted(data["enum_values"]):
                    safe_val = val if not val[0].isdigit() else f"_{val}"
                    lines.append(f"    {safe_val} = '{val}'")
                return lines

            parent = data["parent"]
            base = parent if parent in type_map or parent in self.all_classes else "CIMBase"
            lines.append(f"class {name}({base}):")
            lines.append(f"    __cim_class__ = URIRef('{data['uri']}')")

            prop_names = set(["mRID"])
            props = sorted(data["properties"].values(), key=lambda x: x["name"])
            if not props:
                lines.append("    pass")
            for prop in props:
                p_name = prop["name"]
                if p_name == "mRID":
                    continue
                if p_name in prop_names:
                    p_name = f"{p_name}_{prop['uri'].split('#')[-1].replace('.', '_')}"
                prop_names.add(p_name)

                p_range = prop["range"]
                p_type = type_map.get(p_range, f"'{p_range}'")

                if prop["is_list"]:
                    lines.append(f"    {p_name}: List[{p_type}] = Field(default_factory=list, alias='{prop['uri']}')")
                else:
                    lines.append(f"    {p_name}: Optional[{p_type}] = Field(default=None, alias='{prop['uri']}')")
            lines.append("")
            lines.append(f"CIMModelRegistry.register(URIRef('{data['uri']}'), {name})")
            return lines

        # Sort classes by inheritance before distributing definitions across profile modules.
        defined = set(type_map.keys())
        pending = list(self.all_classes.keys())
        ordered_classes = []
        
        while pending:
            changed = False
            for name in list(pending):
                data = self.all_classes[name]
                parent = data["parent"]
                if parent and parent not in defined and parent in self.all_classes:
                    continue

                ordered_classes.append(name)
                defined.add(name)
                pending.remove(name)
                changed = True
            
            if not changed and pending:
                # Cycle or missing parent
                for name in pending:
                    self.all_classes[name]["parent"] = None

        profile_order = list(self.profiles.keys())
        profile_classes = {profile: [] for profile in profile_order}
        owner = {}
        for name in ordered_classes:
            profile = self.all_classes[name]["profile"]
            if profile not in profile_classes:
                profile_classes[profile] = []
                profile_order.append(profile)
            owner[name] = profile
            profile_classes[profile].append(name)

        builtin_bases = {"CIMBase", "str", "Enum", "object"}
        for profile in profile_order:
            owned = profile_classes[profile]
            imports = {}
            for name in owned:
                parent = self.all_classes[name]["parent"]
                if not parent or parent in builtin_bases or parent in type_map:
                    continue
                parent_owner = owner.get(parent)
                if parent_owner and parent_owner != profile:
                    imports.setdefault(parent_owner, set()).add(parent)

            p_lines = list(module_header)
            for source_profile in profile_order:
                imported_names = sorted(imports.get(source_profile, []))
                if imported_names:
                    p_lines.append(f"from .{source_profile} import {', '.join(imported_names)}")
            if imports:
                p_lines.append("")

            for cls_name in owned:
                p_lines.extend(class_block(cls_name))
                p_lines.append("")

            p_lines.append("__all__ = [")
            for cls_name in owned:
                p_lines.append(f"    '{cls_name}',")
            p_lines.append("]")

            with open(os.path.join(models_dir, f"{profile}.py"), "w") as f:
                f.write("\n".join(p_lines) + "\n")

        aggregate_lines = []
        all_class_names = []
        for profile in profile_order:
            aggregate_lines.append(f"from .{profile} import *")
            all_class_names.extend(profile_classes[profile])
        aggregate_lines.append("")
        aggregate_lines.append("__all__ = [")
        for cls_name in all_class_names:
            aggregate_lines.append(f"    '{cls_name}',")
        aggregate_lines.append("]")
        aggregate_lines.append("")
        aggregate_lines.append("_types_namespace = globals()")
        aggregate_lines.append("for _model_name in __all__:")
        aggregate_lines.append("    _model = _types_namespace[_model_name]")
        aggregate_lines.append("    if hasattr(_model, 'model_rebuild'):")
        aggregate_lines.append("        _model.model_rebuild(_types_namespace=_types_namespace)")
        aggregate_lines.append("")
        aggregate_lines.append("del _model_name, _model, _types_namespace")

        with open(os.path.join(models_dir, "__init__.py"), "w") as f:
            f.write("\n".join(aggregate_lines) + "\n")

        all_lines = ["# Compatibility aggregate for callers that import aceiot_cim_rdf.cim_pydantic.models.all.", ""]
        all_lines.extend(aggregate_lines[: len(profile_order)])
        all_lines.append("")
        all_lines.append("__all__ = [")
        for cls_name in all_class_names:
            all_lines.append(f"    '{cls_name}',")
        all_lines.append("]")
        with open(os.path.join(models_dir, "all.py"), "w") as f:
            f.write("\n".join(all_lines) + "\n")

        # Update __init__.py
        with open(os.path.join(out_dir, "__init__.py"), "w") as f:
            f.write("from .base import CIMBase, CIMModelRegistry, serialize_to_turtle, deserialize_from_turtle\n")
            f.write("from . import models\n")

if __name__ == "__main__":
    parser = SchemaParser("/path/to/ENTSOE_CGMES_v2.4.15_16Feb2016_RDFS")
    parser.parse()
    parser.generate("cim_pydantic")
