import ast
import os
import textwrap

SOURCE_FILE = "main.py"
MODELS_DIR = "../models"

os.makedirs(MODELS_DIR, exist_ok=True)

with open(SOURCE_FILE, "r", encoding="utf-8") as f:
    source_code = f.read()

tree = ast.parse(source_code)

pydantic_classes = []
other_nodes = []

for node in tree.body:
    if isinstance(node, ast.ClassDef):
        if any(
            isinstance(base, ast.Name) and base.id == "BaseModel" for base in node.bases
        ):
            pydantic_classes.append(node)
            continue
    other_nodes.append(node)

imports = """
from pydantic import BaseModel, Field
from typing import Optional, List
"""

for cls in pydantic_classes:
    class_name = cls.name
    file_name = f"{class_name.lower()}.py"
    file_path = os.path.join(MODELS_DIR, file_name)

    class_code = ast.get_source_segment(source_code, cls)
    class_code = textwrap.dedent(class_code)

    with open(file_path, "w", encoding="utf-8") as f:
        f.write(imports.strip() + "\n\n")
        f.write(class_code)

    print(f"✅ Modèle extrait : {class_name} → {file_path}")

new_tree = ast.Module(body=other_nodes, type_ignores=[])
new_source = ast.unparse(new_tree)

with open(SOURCE_FILE, "w", encoding="utf-8") as f:
    f.write(new_source)

print("\n🎉 Extraction terminée")