import os
from llama_cpp import Llama

MODEL_FILENAME = "phi-3-mini-4k-instruct.Q4_K_M.gguf"

current_dir = os.path.dirname(os.path.abspath(__file__))
model_path = os.path.join(current_dir, MODEL_FILENAME)

SYSTEM_PROMPT = """
You are a professional Financial Analyst with expertise in:

- Stock market analysis and valuation
- Corporate finance and accounting
- Investment strategy and portfolio management
- Economic trends and market indicators
- Risk assessment and financial modeling

Your responses should be:

1. Accurate and data-driven
2. Professional and neutral in tone
3. Comprehensive yet concise
4. Based on sound financial principles

Always provide specific examples and metrics when relevant.
"""

# Check if model file exists
if not os.path.exists(model_path):
    print(f"\n❌ Model file '{MODEL_FILENAME}' not found.")
    print("Make sure the GGUF file is in the same folder as main.py\n")
    exit()

print("\n🧠 Loading Financial Analyst AI model...\n")

try:
    llm = Llama(
        model_path=model_path,
        n_ctx=1024,
        n_threads=os.cpu_count(),
        chat_format="phi-3",
        verbose=False
    )

    print("✅ Model loaded successfully!")

except Exception as e:
    print(f"\n❌ Failed to load model:\n{e}")
    exit()

print("\n" + "=" * 60)
print("📈 Financial Analyst AI")
print("Type 'exit' or 'quit' to stop.")
print("=" * 60)

while True:

    user_input = input("\nYou: ")

    if user_input.lower() in ["exit", "quit", "q"]:
        print("\nGoodbye!")
        break

    if not user_input.strip():
        continue

    print("\nAnalyst: ", end="", flush=True)

    try:
        response = llm.create_chat_completion(
            messages=[
                {
                    "role": "system",
                    "content": SYSTEM_PROMPT
                },
                {
                    "role": "user",
                    "content": user_input
                }
            ],

            max_tokens=512,
            temperature=0.2,
            top_p=0.9,

            stop=[
                "<|user|>",
                "<|assistant|>"
            ],

            stream=True
        )

        for chunk in response:
            delta = chunk["choices"][0]["delta"]

            if "content" in delta:
                print(delta["content"], end="", flush=True)

        print()

    except KeyboardInterrupt:
        print("\n\nStopped by user.")
        break

    except Exception as e:
        print(f"\n❌ Error: {e}")