import requests
import time
import json
import asyncio
import os
from chatpbc_file_processor import process_file
from chatpbc_webscraper import get_business_intelligence_brief

HF_TOKEN = "YOUR_HF_TOKEN_HERE"
ROUTER_ENDPOINT = "https://router.huggingface.co/v1/chat/completions"

CHAT_TEMPLATES = {
    "chatpbc-v4": {
        "model_id": "chatpbc1/chatpbc-v4",
        "system_prompt": "You are ChatPBC V4, the apex AI business strategist and intelligence analyst developed by Mik Tse Agency. You have access to real-time website data and uploaded business documents provided by the user. Your role is to deliver world-class business consulting: strategic analysis, competitive intelligence, market research, financial modeling guidance, M&A advisory, go-to-market strategy, operational efficiency, and organizational transformation. You cover 26 industries: Technology, Finance, Healthcare, Retail, Manufacturing, Energy, Telecom, Automotive, Real Estate, Media, Travel, Food & Beverage, Agriculture, Education, Government, Consulting, Logistics, Marketing, Human Resources, Legal, Non-profit, Biotechnology, Aerospace & Defense, Fashion, Sports & Entertainment, and Environmental Services. When given website data or files, analyze them deeply and provide actionable strategic insights. Always maintain full conversation memory. Respond with the depth and precision of a McKinsey senior partner."
    },
    "chatpbc-v33": {
        "model_id": "chatpbc1/chatpbc-v33",
        "system_prompt": "You are ChatPBC V3.3, a highly intelligent and conversational AI business consultant developed by Mik Tse Agency. You have access to real-time website data and uploaded business documents provided by the user. You are warm, professional, and strategic. You respond like a real human consultant: you greet users, ask follow-up questions, show empathy when businesses are struggling, and provide clear, actionable advice. You cover 26 industries: Technology, Finance, Healthcare, Retail, Manufacturing, Energy, Telecom, Automotive, Real Estate, Media, Travel, Food & Beverage, Agriculture, Education, Government, Consulting, Logistics, Marketing, Human Resources, Legal, Non-profit, Biotechnology, Aerospace & Defense, Fashion, Sports & Entertainment, and Environmental Services. When given website data or files, analyze them and provide practical, implementable recommendations. Always maintain full conversation memory."
    }
}

async def query_model(model_id, user_message, conversation_history=None, files=None, url=None):
    model_config = CHAT_TEMPLATES.get(model_id)
    if not model_config:
        raise ValueError(f"Model ID {model_id} not found in CHAT_TEMPLATES.")

    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }

    context_str = ""
    if files:
        for file_path in files:
            processed_content = process_file(file_path)
            context_str += f"=== File: {file_path.split('/')[-1]} ===\n{processed_content}\n\n"

    if url:
        print(f"Scraping URL: {url}")
        web_brief = await get_business_intelligence_brief(url)
        context_str += f"=== Website Analysis: {url} ===\n{web_brief}\n\n"

    full_user_message = f"{context_str}{user_message}" if context_str else user_message

    messages = [{"role": "system", "content": model_config["system_prompt"]}]
    if conversation_history:
        for turn in conversation_history:
            messages.append(turn)
    messages.append({"role": "user", "content": full_user_message})

    payload = {
        "model": model_config["model_id"],
        "messages": messages,
        "max_tokens": 1024,
        "temperature": 0.7,
        "stream": False
    }

    for i in range(5): # 5 retries
        try:
            response = requests.post(ROUTER_ENDPOINT, headers=headers, json=payload, timeout=300)
            if response.status_code in [503, 429]:
                print(f"Model loading or rate limited, retrying in 10 seconds... (Attempt {i+1}/5)")
                time.sleep(10)
                continue
            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"]
        except requests.exceptions.RequestException as e:
            print(f"Attempt {i+1} failed: {e}")
            if i == 4:
                raise Exception("Failed to get response after multiple retries.")
            time.sleep(10)

async def main():
    # Example usage
    print("Querying ChatPBC V4...")
    try:
        response_v4 = await query_model(
            "chatpbc-v4",
            "What are the key business insights for a new tech startup?",
            files=[],
            url=None
        )
        print("ChatPBC V4 Response:", response_v4)
    except Exception as e:
        print(f"Error with ChatPBC V4: {e}")

if __name__ == "__main__":
    asyncio.run(main())