import argparse
import json
from src.ontology.matcher import ConceptMatcher
from src.parser.parser import Parser
from src.enrichment.enricher import Enricher
from src.prompt_builder.builder import PromptBuilder

from src.embeddings.engine import EmbeddingEngine

def main():
    parser = argparse.ArgumentParser(description="Prompt Compiler CLI")
    parser.add_argument("text", help="Natural language description")
    parser.add_argument("--ontology", default="data/ontology", help="Path to ontology directory")
    parser.add_argument("--index-dir", default="data/faiss_indices", help="Path to FAISS indices directory")
    parser.add_argument("--json", action="store_true", help="Output IR as JSON")
    parser.add_argument("--safe", action="store_true", default=True, help="Enable safe mode")
    parser.add_argument("--no-safe", action="store_false", dest="safe", help="Disable safe mode")
    parser.add_argument("--semantic", action="store_true", help="Enable semantic retrieval")
    parser.add_argument("--profile", default="generic", choices=["generic", "sdxl", "pony", "illustrious"], help="Prompt profile to use")
    parser.add_argument("--runtime", default="pytorch", choices=["pytorch", "onnx"], help="Inference runtime to use")
    parser.add_argument("--language", default="en", choices=["en", "es"], help="Input language")

    args = parser.parse_args()

    # Initialize components
    matcher = ConceptMatcher(args.ontology)
    
    embedding_engine = None
    if args.semantic:
        try:
            if args.runtime == "onnx":
                from src.runtime.onnx_runtime import ONNXEmbeddingEngine
                embedding_engine = ONNXEmbeddingEngine(index_dir=args.index_dir + "_onnx")
            else:
                from src.embeddings.engine import EmbeddingEngine
                embedding_engine = EmbeddingEngine(index_dir=args.index_dir)
            embedding_engine.load_index()
        except FileNotFoundError:
            print("Warning: FAISS indices not found. Semantic retrieval disabled.")
            embedding_engine = None
            
    parser_engine = Parser(matcher, embedding_engine=embedding_engine, language=args.language)
    enricher = Enricher(matcher)
    builder = PromptBuilder(profile_name=args.profile)

    # Run pipeline
    ir = parser_engine.parse(args.text, safe_mode=args.safe)
    ir = enricher.enrich(ir)
    prompt_bundle = builder.build(ir)

    if args.json:
        output = {
            "ir": ir.model_dump(),
            "prompt": prompt_bundle
        }
        print(json.dumps(output, indent=2))
    else:
        print(f"\n--- Explainability Trace ---")
        print(json.dumps(ir.trace, indent=2))
        print(f"\n--- Positive Prompt ---")
        print(prompt_bundle["positive"])
        print(f"\n--- Negative Prompt ---")
        print(prompt_bundle["negative"])

if __name__ == "__main__":
    main()