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"

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            target = None
            if "Global" in node.name: target = "global"
            elif "Domain" in node.name: target = "domain"
            
            if 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)
    global_fields = ctx_data.get('global', {})
    domain_fields = ctx_data.get('domain', {})

    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("")

    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]' # TODO: Support typed lists
        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
    generate_class("AppGlobal", "BaseGlobalContext", global_fields)
    generate_class("AppDomain", "BaseDomainContext", domain_fields)

    lines.append("@dataclass")
    lines.append("class SystemContext(BaseSystemContext):")
    lines.append("    pass")

    return "\n".join(lines)
