import os
import json
import asyncio
import logging
from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.responses import StreamingResponse
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from tavily import TavilyClient
from groq import Groq

logging.basicConfig(level=logging.INFO, format='%(message)s')

groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))

logging.info("Downloading Konkani GGUF Model from friend's repo...")
model_path = hf_hub_download(
    repo_id="dom0804/konkani_companion_GGUF", 
    filename="qwen2.5-3b-instruct.Q4_K_M.gguf"               
)

logging.info("Loading into CPU...")
llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2)

app = FastAPI()

class ChatPayload(BaseModel):
    messages: list
    stream: bool = False

@app.post("/openai/v1/chat/completions")
async def chat_endpoint(payload: ChatPayload):
    user_query = payload.messages[-1]["content"]

    # --- PHASE 1: SMART REWRITER ---
    history_text = "\n".join([f"{m['role']}: {m['content']}" for m in payload.messages[:-1]])
    rewriter_instruction = f"""You are a Search Query Optimizer.
    RULES:
    1. If the query is purely conversational (e.g., "hi", "how are you"), output EXACTLY: SKIP_SEARCH
    2. If it is factual, generate a SINGLE optimized Google search string. 
    3. Output ONLY the raw search query or SKIP_SEARCH.
    History: {history_text}
    Latest Query: {user_query}
    """
    
    try:
        rewriter_response = groq_client.chat.completions.create(
            messages=[{"role": "user", "content": rewriter_instruction}],
            model="llama-3.1-8b-instant",
            temperature=0.1
        )
        optimized_query = rewriter_response.choices[0].message.content.strip()
    except Exception:
        optimized_query = user_query 

    # --- PHASE 2: SMART SEARCH ---
    if "SKIP_SEARCH" in optimized_query:
        fact = "This is a conversational query. Respond naturally."
    else:
        try:
            search_data = tavily_client.search(query=optimized_query, search_depth="basic")
            fact = "\n".join([result['content'] for result in search_data['results'][:2]])
        except Exception:
            fact = "No live info available."

    # --- PHASE 3: GROQ BRAIN ---
    brain_instruction = f"""<system_role>
    You are the Core Reasoning Engine for a cross-cultural Konkani conversational AI. 
    </system_role>
    <live_context>{fact}</live_context>
    <operational_rules>
    1. CONVERSATIONAL FLUIDITY: Respond warmly and naturally to social queries.
    2. FACTUAL GROUNDING: Base answers on the <live_context>. Do not hallucinate.
    3. DIRECT TRANSLATION OVERRIDE: If asked to translate, output EXACTLY the target English text.
    4. DOWNSTREAM SAFETY: Max 3 sentences per paragraph. NO markdown.
    </operational_rules>"""

    groq_messages = [{"role": "system", "content": brain_instruction}]
    for msg in payload.messages[:-1]:
        groq_messages.append({"role": msg["role"], "content": msg["content"]})
    groq_messages.append({"role": "user", "content": user_query})

    response = groq_client.chat.completions.create(
        messages=groq_messages,
        model="llama-3.3-70b-versatile"
    )
    english_paragraphs = [p for p in response.choices[0].message.content.strip().split('\n') if p.strip()]

    # --- PHASE 4: GGUF TRANSLATION STREAM ---
    async def stream_generator():
        for para in english_paragraphs:
            formatted_prompt = f"<|im_start|>system\nYou are a highly accurate translation model. Translate the following English text into Konkani.<|im_end|>\n<|im_start|>user\n{para}<|im_end|>\n<|im_start|>assistant\n"
            
            stream = llm(formatted_prompt, max_tokens=1024, stop=["<|im_end|>"], stream=True, temperature=0.1)
            for chunk in stream:
                yield f"data: {json.dumps({'id': 'chatcmpl-custom', 'object': 'chat.completion.chunk', 'choices': [{'delta': {'content': chunk['choices'][0]['text']}}]})}\n\n"
                await asyncio.sleep(0.01)
            yield f"data: {json.dumps({'choices': [{'delta': {'content': '\n\n'}}]})}\n\n"
        yield "data: [DONE]\n\n"

    if payload.stream:
        return StreamingResponse(stream_generator(), media_type="text/event-stream")
    else:
        final_text = ""
        async for chunk in stream_generator():
            if chunk != "data: [DONE]\n\n":
                data_dict = json.loads(chunk.replace("data: ", "").strip())
                if "content" in data_dict["choices"][0]["delta"]:
                    final_text += data_dict["choices"][0]["delta"]["content"]
        return {"choices": [{"message": {"role": "assistant", "content": final_text.strip()}}]}