import yaml
import ast
from pathlib import Path
from typing import Dict, Any, List


def generate_schema_from_file(context_path: str) -> Dict[str, Any]:
    """
    Parses a python context file (e.g. src/context.py) and generates a YAML schema representation.
    Used for 'schema gen' command.
    """
    path = Path(context_path)
    if not path.exists():
        raise FileNotFoundError(f"Context file not found: {context_path}")

    with open(path, "r", encoding="utf-8") as f:
        tree = ast.parse(f.read())

    schema = {"context": {"global": {}, "domain": {}}}

    # Helper to parse AST assignment
    def parse_annotation(ann):
        if isinstance(ann, ast.Name):
            return ann.id
        elif isinstance(ann, ast.Subscript):  # e.g. List[str]
            val = parse_annotation(ann.slice)
            return f"{ann.value.id}[{val}]"
        return "Any"

    # Parse all classes found in the file
    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            # Resolve target namespace name (camelCase ClassName -> lower_snake_case or just keyword)
            node_name_lower = node.name.lower()
            target = None
            
            if "global" in node_name_lower:
                target = "global"
            elif "domain" in node_name_lower:
                target = "domain"
            elif node_name_lower not in ("systemcontext", "basesystemcontext"):
                # Generic namespace detection (experimental)
                target = node_name_lower.replace("context", "")
            
            if target:
                if target not in schema["context"]:
                    schema["context"][target] = {}
                    
                for item in node.body:
                    if isinstance(item, ast.AnnAssign) and isinstance(
                        item.target, ast.Name
                    ):
                        field_name = item.target.id
                        field_type = parse_annotation(item.annotation)

                        # Simplified mapping back to YAML types
                        yaml_type = "string"
                        if "int" in field_type:
                            yaml_type = "integer"
                        elif "float" in field_type:
                            yaml_type = "float"
                        elif "bool" in field_type:
                            yaml_type = "boolean"
                        elif "List" in field_type:
                            yaml_type = "list"
                        elif "Dict" in field_type:
                            yaml_type = "dict"

                        field_def = {"type": yaml_type}

                        # Detect default value
                        if item.value:
                            try:
                                # Simple literals
                                val = ast.literal_eval(item.value)
                                field_def["default"] = val
                            except:
                                pass

                        schema["context"][target][field_name] = field_def

    return schema


def generate_code_from_schema(schema_path: str) -> str:
    """
    Generates Python code (dataclasses) from a YAML context schema.
    Used for 'schema code' command.
    """
    path = Path(schema_path)
    if not path.exists():
        raise FileNotFoundError(f"Schema file not found: {schema_path}")

    with open(path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f) or {}

    # Handle root 'context' key
    ctx_data = data.get("context", data)

    lines = []
    lines.append("# Generated by Theus Context Code Gen")
    lines.append("from dataclasses import dataclass, field")
    lines.append("from typing import Optional, List, Dict, Any")
    lines.append(
        "from theus import BaseSystemContext, BaseGlobalContext, BaseDomainContext"
    )
    lines.append("")
    
    # ... (map_type and generate_class remain same or updated) ...
    def map_type(t: str) -> str:
        t = t.lower()
        if t == "string":
            return "str"
        if t == "integer":
            return "int"
        if t == "boolean":
            return "bool"
        if t == "float":
            return "float"
        if t == "list":
            return "List[Any]"
        if t == "dict":
            return "Dict[str, Any]"
        return "Any"

    def generate_class(name: str, parent: str, fields: Dict):
        lines.append("@dataclass")
        lines.append(f"class {name}({parent}):")
        if not fields:
            lines.append("    pass")
            lines.append("")
            return

        for fname, fspec in fields.items():
            ftype = map_type(fspec.get("type", "string"))
            default = fspec.get("default")

            line = f"    {fname}: {ftype}"
            if default is not None:
                if isinstance(default, str):
                    line += f" = '{default}'"
                else:
                    line += f" = {default}"
            elif ftype.startswith("List") or ftype.startswith("Dict"):
                line += (
                    " = field(default_factory=list)"
                    if ftype.startswith("List")
                    else " = field(default_factory=dict)"
                )

            lines.append(line)
        lines.append("")

    # Generate Classes for each namespace
    for ns_name, fields in ctx_data.items():
        if ns_name == "global":
            generate_class("AppGlobal", "BaseGlobalContext", fields)
        elif ns_name == "domain":
            generate_class("AppDomain", "BaseDomainContext", fields)
        else:
            # Capitalize name for Class
            class_name = ns_name.capitalize() + "Context"
            generate_class(class_name, "BaseDomainContext", fields)

    lines.append("@dataclass")
    lines.append("class SystemContext(BaseSystemContext):")
    # Add explicit fields for non-standard namespaces if needed
    for ns_name in ctx_data:
        if ns_name not in ("global", "domain"):
             class_name = ns_name.capitalize() + "Context"
             lines.append(f"    {ns_name}: {class_name}")
    
    if not any(k not in ("global", "domain") for k in ctx_data):
        lines.append("    pass")

    return "\n".join(lines)
