from typing import Any, Dict, List, Optional, Type, Union
from pydantic import BaseModel, Field, ConfigDict
from rdflib import Graph, Literal, Namespace, URIRef, RDF, RDFS
import uuid

CIM = Namespace("http://iec.ch/TC57/2013/CIM-schema-cim16#")
DCTERMS = Namespace("http://purl.org/dc/terms/")

class CIMBase(BaseModel):
    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)

    mRID: str = Field(default_factory=lambda: str(uuid.uuid4()), alias="mRID")
    
    def to_rdf(self, graph: Graph):
        subject = URIRef(f"#{self.mRID}")
        cls_uri = getattr(self, "__cim_class__", None)
        if cls_uri:
            graph.add((subject, RDF.type, cls_uri))
        
        for field_name, field_info in self.model_fields.items():
            if field_name == "mRID":
                continue
            
            value = getattr(self, field_name)
            if value is None:
                continue
            
            prop_uri = field_info.alias
            if not prop_uri:
                continue
            
            if isinstance(value, list):
                for v in value:
                    self._add_to_graph(graph, subject, prop_uri, v)
            else:
                self._add_to_graph(graph, subject, prop_uri, value)

    def _add_to_graph(self, graph, subject, prop_uri, value):
        prop = URIRef(prop_uri)
        
        if isinstance(value, CIMBase):
            graph.add((subject, prop, URIRef(f"#{value.mRID}")))
        elif isinstance(value, str):
            # If it starts with #, or if it doesn't contain spaces and looks like an ID
            # A better way: if it matches any known mRID in the current export.
            # For now, let's just check for # or if it's a common prefix we use.
            # Actually, let's just assume if it's a string and the property is an object property, it's an mRID.
            # But we don't know if it's an object property here easily.
            is_ref = value.startswith("#") or "_" in value or any(value.startswith(p) for p in ["sub", "vl", "cn", "tn", "load", "gen", "line", "term", "bv", "gu", "trans", "pte"])
            if is_ref:
                mrid = value.lstrip("#")
                graph.add((subject, prop, URIRef(f"#{mrid}")))
            else:
                graph.add((subject, prop, Literal(value)))
        elif isinstance(value, (int, float, bool)):
            graph.add((subject, prop, Literal(value)))
        elif isinstance(value, URIRef):
            graph.add((subject, prop, value))

    @classmethod
    def from_rdf(cls, graph: Graph, subject: URIRef):
        subject_str = str(subject)
        mrid = subject_str.split("#")[-1] if "#" in subject_str else subject_str
        data = {"mRID": mrid}
        
        for _, p, o in graph.triples((subject, None, None)):
            if p == RDF.type:
                continue
            
            field_name = None
            for name, info in cls.model_fields.items():
                if info.alias == str(p):
                    field_name = name
                    break
            
            if not field_name:
                continue
            
            value = None
            if isinstance(o, URIRef):
                o_str = str(o)
                mrid = o_str.split("#")[-1] if "#" in o_str else (o_str.split(".")[-1] if "." in o_str else o_str)
                
                # Check if this field expects a CIMBase or List[CIMBase]
                # We need to get the actual type of the field
                field_type = info.annotation
                
                # Handle Optional and List
                from typing import get_args, get_origin
                origin = get_origin(field_type)
                if origin is Union:
                    args = get_args(field_type)
                    # Use the first non-None type
                    field_type = next(a for a in args if a is not type(None))
                
                # Re-check origin for List
                origin = get_origin(field_type)
                if origin is list:
                    field_type = get_args(field_type)[0]

                # Now field_type should be the base class (e.g. Simple_Float or ConductingEquipment)
                # If it's a string, we might have a forward ref or it might be a real string field.
                # In our generated models, they are often strings 'Simple_Float' etc.
                if isinstance(field_type, str):
                    # Try to resolve from registry by name? 
                    # For now, let's just see if it's a CIMBase reference in the graph
                    pass

                # If the object 'o' has a type in the graph, and it's a CIMBase, let's try to instantiate it
                types = list(graph.objects(o, RDF.type))
                if types:
                    t = types[0]
                    model_cls = CIMModelRegistry.get(t)
                    if model_cls:
                        # Recursive call
                        value = model_cls.from_rdf(graph, o)
                    else:
                        value = mrid
                else:
                    value = mrid
            elif isinstance(o, Literal):
                value = o.toPython()
            
            if value is not None:
                if field_name in data:
                    if not isinstance(data[field_name], list):
                        data[field_name] = [data[field_name]]
                    data[field_name].append(value)
                else:
                    data[field_name] = value
        
        return cls(**data)

def serialize_to_turtle(objects: List[CIMBase], metadata: Optional[Dict[str, Any]] = None) -> str:
    g = Graph()
    g.bind("cim", CIM)
    g.bind("dcterms", DCTERMS)
    for obj in objects:
        obj.to_rdf(g)
    if metadata:
        _add_metadata(g, metadata)
    return g.serialize(format="turtle")

def _add_metadata(graph: Graph, metadata: Dict[str, Any]):
    subject = URIRef(metadata.get("subject", "#dataset"))
    for prop_uri, value in metadata.get("properties", {}).items():
        values = value if isinstance(value, list) else [value]
        for item in values:
            if isinstance(item, URIRef):
                graph.add((subject, URIRef(prop_uri), item))
            else:
                graph.add((subject, URIRef(prop_uri), Literal(item)))

def deserialize_from_turtle(data: str) -> List[CIMBase]:
    g = Graph()
    g.parse(data=data, format="turtle")
    
    objects = []
    # Find all subjects with an RDF type
    subjects = set(g.subjects(RDF.type, None))
    for s in subjects:
        types = list(g.objects(s, RDF.type))
        for t in types:
            model_cls = CIMModelRegistry.get(t)
            if model_cls:
                objects.append(model_cls.from_rdf(g, s))
                break
    return objects

class CIMModelRegistry:
    _registry: Dict[URIRef, Type[CIMBase]] = {}

    @classmethod
    def register(cls, uri: URIRef, model: Type[CIMBase]):
        cls._registry[uri] = model

    @classmethod
    def get(cls, uri: URIRef) -> Optional[Type[CIMBase]]:
        return cls._registry.get(uri)
