import os
from pathlib import Path

# Mapping of old import prefixes to new import prefixes
IMPORT_MAP = {
    "app.services.ai.chat_engine": "app.services.chat.chat_engine",
    "app.services.ai.memory_manager": "app.services.chat.memory_manager",
    "app.services.ai.prompts": "app.services.chat.prompts",
    "app.services.ai.faq_data": "app.services.chat.faq_data",
    "app.services.ai.faq_matcher": "app.services.chat.faq_matcher",
    
    "app.services.ai.prediction_engine": "app.services.prediction.prediction_engine",
    "app.ml_artifacts.trainer": "app.services.prediction.trainer.trainer",
    
    "app.services.ai.function_registry": "app.services.tools.function_registry",
    "app.services.ai.tool_executor": "app.services.tools.tool_executor",
    
    "app.services.ai.document_analyzer": "app.services.document.document_analyzer",
    
    "app.services.stt": "app.services.voice",
    "app.services.ml": "app.services.feedback",
    
    "app.services.ai": "app.services", # fallback just in case
}

# The joblib file paths in code needs updating
PATH_MAP = {
    '"app/services/prediction/trainer/models/': '"app/services/prediction/trainer/models/',
    "'app/services/prediction/trainer/models/": "'app/services/prediction/trainer/models/",
    '"app/services/ocr/providers/': '"app/services/ocr/providers/providers/' # this isn't stringly typed, but let's see
}

TARGET_DIRS = ["app", "scripts"]

def process_file(filepath: Path):
    content = filepath.read_text(encoding="utf-8")
    original = content
    
    for old, new in IMPORT_MAP.items():
        content = content.replace(f"from {old}", f"from {new}")
        content = content.replace(f"import {old}", f"import {new}")
    
    for old, new in PATH_MAP.items():
         content = content.replace(old, new)
         
    # Fix OCR providers paths inside ocr_router.py
    if filepath.name == "ocr_router.py":
        content = content.replace("app.services.ocr.groq_provider", "app.services.ocr.providers.groq_provider")
        content = content.replace("app.services.ocr.gemini_provider", "app.services.ocr.providers.gemini_provider")
        content = content.replace("app.services.ocr.openai_provider", "app.services.ocr.providers.openai_provider")
        content = content.replace("app.services.ocr.huggingface_provider", "app.services.ocr.providers.huggingface_provider")

    if content != original:
        filepath.write_text(content, encoding="utf-8")
        print(f"Updated {filepath}")

for d in TARGET_DIRS:
    for root, _, files in os.walk(d):
        for f in files:
            if f.endswith(".py") or f.endswith(".md"):
                process_file(Path(root) / f)

# Also update main.py in root
process_file(Path("main.py"))

print("Refactoring complete.")