app/main.py

import os
from transformers import safetensors_conversion
safetensors_conversion.auto_conversion = lambda *args, **kwargs: None
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"

import uuid
import base64
import logging
from huggingface_hub import login, get_token
from fastapi import FastAPI, UploadFile, File, Header, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, Literal


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("farmlingua")

hf_token = os.environ.get("HF_TOKEN") or get_token()
if hf_token:
    login(token=hf_token)
else:
    raise RuntimeError("HF_TOKEN not found.")

from app.memory import get_history, append_turn, clear_session
from app.text_to_text.farm_agent import farm_agent
from app.speech_to_text.speech_agent import speech_agent

app = FastAPI(title="FarmLingua AI", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=[
        "X-UID", "X-Transcription", "X-Language",
        "X-Confidence", "X-English-Input",
    ],
)



def resolve_uid(x_uid: Optional[str]) -> str:
    return x_uid if x_uid else str(uuid.uuid4())


def encode_header(value: str) -> str:
    return base64.b64encode(value.encode("utf-8")).decode("ascii")


def is_sentence_boundary(text: str) -> bool:
    stripped = text.strip()
    return (
        stripped.endswith((".", "!", "?", "\n")) and
        len(stripped) > 10
    )


def stream_with_translation(
    uid: str,
    channel: str,
    english_input: str,
    detected_lang: str,
    history: list,
):
    streamer       = farm_agent.stream_response(history, english_input)
    english_answer = ""
    buffer         = ""
    chunk_count    = 0

    for token in streamer:
        english_answer += token
        buffer         += token

        if is_sentence_boundary(buffer):
            chunk = farm_agent._clean_llm_output(buffer.strip())
            buffer = ""

            if not chunk:
                continue

            chunk_count += 1

            if detected_lang != "english":
                log.info(
                    f"[TRANSLATE BACK] chunk {chunk_count} | "
                    f"english → {detected_lang} | "
                    f"EN: {chunk[:80]}..."
                )
                translated = farm_agent.translate(
                    chunk,
                    src_lang="english",
                    tgt_lang=detected_lang,
                )
                log.info(
                    f"[TRANSLATE BACK] chunk {chunk_count} | "
                    f"result: {translated[:80]}..."
                )
                yield translated + " "
            else:
                yield chunk + " "

    # Flush remaining buffer
    if buffer.strip():
        chunk = farm_agent._clean_llm_output(buffer.strip())
        if chunk:
            chunk_count += 1
            if detected_lang != "english":
                log.info(
                    f"[TRANSLATE BACK] flush chunk {chunk_count} | "
                    f"english → {detected_lang} | "
                    f"EN: {chunk[:80]}..."
                )
                translated = farm_agent.translate(
                    chunk,
                    src_lang="english",
                    tgt_lang=detected_lang,
                )
                log.info(
                    f"[TRANSLATE BACK] flush result: {translated[:80]}..."
                )
                yield translated
            else:
                yield chunk

    log.info(
        f"[PIPELINE DONE] uid={uid} | channel={channel} | "
        f"total chunks={chunk_count} | "
        f"english answer preview: {english_answer[:120]}..."
    )

    append_turn(uid, channel, "assistant", english_answer.strip())


def stream_text_pipeline(uid: str, channel: str, user_text: str):
    meta          = farm_agent.process(user_text, get_history(uid, channel))
    detected_lang = meta["detected_lang"]
    confidence    = meta["confidence"]
    english_input = meta["english_input"]

    log.info(f"[TEXT PIPELINE] uid={uid}")
    log.info(f"  Original text   : {user_text[:120]}")
    log.info(f"  Detected lang   : {detected_lang} ({confidence:.2%} confidence)")
    log.info(f"  English input   : {english_input[:120]}")

    append_turn(uid, channel, "user", english_input)
    history = get_history(uid, channel)[:-1]

    yield from stream_with_translation(
        uid, channel, english_input, detected_lang, history
    )


def stream_stt_pipeline(uid: str, channel: str, transcription: str, language: str):
    log.info(f"[STT PIPELINE] uid={uid}")
    log.info(f"  Selected lang   : {language}")
    log.info(f"  Transcription   : {transcription[:120]}")

    if language != "english":
        english_input = speech_agent.translate_to_english(transcription, language)
        log.info(f"  English input   : {english_input[:120]}")
    else:
        english_input = transcription
        log.info(f"  English input   : (no translation needed)")

    append_turn(uid, channel, "user", english_input)
    history = get_history(uid, channel)[:-1]

    yield from stream_with_translation(
        uid, channel, english_input, language, history
    )




class TextRequest(BaseModel):
    message: str


@app.post("/text/chat")
async def text_chat(
    body: TextRequest,
    x_uid: Optional[str] = Header(default=None),
):
    message = body.message.strip()
    if not message:
        raise HTTPException(status_code=400, detail="Message cannot be empty.")

    uid     = resolve_uid(x_uid)
    headers = {
        "X-UID": uid,
        "Access-Control-Expose-Headers": "X-UID",
    }

    log.info(f"[REQUEST] /text/chat | uid={uid} | message={message[:80]}")

    return StreamingResponse(
        stream_text_pipeline(uid, "text", message),
        media_type="text/plain",
        headers=headers,
    )


@app.post("/stt/chat")
async def stt_chat(
    audio: UploadFile = File(...),
    language: Literal["yoruba", "igbo", "hausa", "english"] = "english",
    x_uid: Optional[str] = Header(default=None),
):
    uid         = resolve_uid(x_uid)
    audio_bytes = await audio.read()

    log.info(f"[REQUEST] /stt/chat | uid={uid} | language={language} | size={len(audio_bytes)/1024:.1f}KB")

    try:
        transcription = speech_agent.transcribe(audio_bytes, language)
        log.info(f"[STT] Transcription: {transcription}")
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))

    headers = {
        "X-UID":           uid,
        "X-Transcription": encode_header(transcription),
        "X-Language":      language,
        "Access-Control-Expose-Headers": "X-UID, X-Transcription, X-Language",
    }

    return StreamingResponse(
        stream_stt_pipeline(uid, "stt", transcription, language),
        media_type="text/plain",
        headers=headers,
    )


@app.delete("/session")
async def clear_user_session(x_uid: str = Header(...)):
    clear_session(x_uid)
    return {"status": "cleared", "uid": x_uid}


@app.get("/health")
async def health():
    return {"status": "ok"}