import json
import os
import time
import uuid
import logging
import asyncio
import concurrent.futures
from typing import AsyncIterator, Optional, Any, Literal, Union

import grpc
import brunix_pb2
import brunix_pb2_grpc

from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("openai-proxy")

_thread_pool = concurrent.futures.ThreadPoolExecutor(
    max_workers=int(os.getenv("PROXY_THREAD_WORKERS", "20"))
)

GRPC_TARGET = os.getenv("BRUNIX_GRPC_TARGET", "localhost:50051")
PROXY_MODEL = os.getenv("PROXY_MODEL_ID", "brunix")

_channel: Optional[grpc.Channel] = None
_stub:    Optional[brunix_pb2_grpc.AssistanceEngineStub] = None


def get_stub() -> brunix_pb2_grpc.AssistanceEngineStub:
    global _channel, _stub
    if _stub is None:
        _channel = grpc.insecure_channel(GRPC_TARGET)
        _stub    = brunix_pb2_grpc.AssistanceEngineStub(_channel)
        logger.info(f"[gRPC] connected to {GRPC_TARGET}")
    return _stub


app = FastAPI(
    title="Brunix OpenAI-Compatible Proxy",
    version="2.0.0",
    description="stream:false → AskAgent | stream:true → AskAgentStream",
)

class ChatMessage(BaseModel):
    role:    Literal["system", "user", "assistant", "function"] = "user"
    content: str = ""
    name:    Optional[str] = None


class ChatCompletionRequest(BaseModel):
    model:             str = PROXY_MODEL
    messages:          list[ChatMessage]
    stream:            bool = False
    temperature:       Optional[float] = None
    max_tokens:        Optional[int]   = None
    session_id:        Optional[str]   = None   # extensión Brunix
    top_p:             Optional[float] = None
    n:                 Optional[int]   = 1
    stop:              Optional[Any]   = None
    presence_penalty:  Optional[float] = None
    frequency_penalty: Optional[float] = None
    user:              Optional[str]   = None


class CompletionRequest(BaseModel):
    model:       str = PROXY_MODEL
    prompt:      Union[str, list[str]] = ""
    stream:      bool = False
    temperature: Optional[float] = None
    max_tokens:  Optional[int]   = None
    session_id:  Optional[str]   = None
    suffix:      Optional[str]   = None
    top_p:       Optional[float] = None
    n:           Optional[int]   = 1
    stop:        Optional[Any]   = None
    user:        Optional[str]   = None


# Ollama schemas
class OllamaChatMessage(BaseModel):
    role:    str = "user"
    content: str = ""


class OllamaChatRequest(BaseModel):
    model:      str = PROXY_MODEL
    messages:   list[OllamaChatMessage]
    stream:     bool = True             # Ollama hace stream por defecto
    session_id: Optional[str] = None


class OllamaGenerateRequest(BaseModel):
    model:      str = PROXY_MODEL
    prompt:     str = ""
    stream:     bool = True
    session_id: Optional[str] = None


def _ts() -> int:
    return int(time.time())


def _chat_response(content: str, req_id: str) -> dict:
    return {
        "id": req_id, "object": "chat.completion", "created": _ts(),
        "model": PROXY_MODEL,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }


def _completion_response(text: str, req_id: str) -> dict:
    return {
        "id": req_id, "object": "text_completion", "created": _ts(),
        "model": PROXY_MODEL,
        "choices": [{"text": text, "index": 0, "logprobs": None, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }


def _chat_chunk(delta: str, req_id: str, finish: Optional[str] = None) -> dict:
    return {
        "id": req_id, "object": "chat.completion.chunk", "created": _ts(),
        "model": PROXY_MODEL,
        "choices": [{"index": 0,
                     "delta": {"role": "assistant", "content": delta} if delta else {},
                     "finish_reason": finish}],
    }


def _completion_chunk(text: str, req_id: str, finish: Optional[str] = None) -> dict:
    return {
        "id": req_id, "object": "text_completion", "created": _ts(),
        "model": PROXY_MODEL,
        "choices": [{"text": text, "index": 0, "logprobs": None, "finish_reason": finish}],
    }


def _sse(data: dict) -> str:
    return f"data: {json.dumps(data)}\n\n"


def _sse_done() -> str:
    return "data: [DONE]\n\n"


def _query_from_messages(messages: list[ChatMessage]) -> str:
    for m in reversed(messages):
        if m.role == "user":
            return m.content
    return ""


async def _invoke_blocking(query: str, session_id: str, context = {}) -> str:

    loop = asyncio.get_event_loop()

    def _call():
        stub = get_stub()


        try:
            ed_contxt = context["editor_content"] or ""
        except Exception:
            ed_contxt = ""

        try:
            sel_contxt = context["selected_text"] or ""
        except Exception:
            sel_contxt = ""

        try:
            ext_contxt = context["extra_context"] or ""
        except Exception:
            ext_contxt = ""

        try:
            us_info = str(context["user_info"]) or "{}"
        except Exception:
            us_info = "{}"


        req  = brunix_pb2.AgentRequest(query=query, session_id=session_id,
                                       editor_content=ed_contxt,
                                       selected_text=sel_contxt,
                                       extra_context=ext_contxt,
                                       user_info=us_info)
                                       
                                       
        parts = []
        for resp in stub.AskAgent(req):
            if resp.text:
                parts.append(resp.text)
        return "".join(parts)

    return await loop.run_in_executor(_thread_pool, _call)


async def _iter_stream(query: str, session_id: str, context = {}) -> AsyncIterator[brunix_pb2.AgentResponse]:

    loop  = asyncio.get_event_loop()
    queue: asyncio.Queue = asyncio.Queue()

    def _producer():
        try:
            stub = get_stub()
            print("CONTEXT ====")
            print(context)
            print("======= ====")

            try:
                ed_contxt = context["editor_content"] or ""
            except Exception:
                ed_contxt = ""

            try:
                sel_contxt = context["selected_text"] or ""
            except Exception:
                sel_contxt = ""

            try:
                ext_contxt = context["extra_context"] or ""
            except Exception:
                ext_contxt = ""

            try:
                us_info = str(context["user_info"]) or "{}"
            except Exception:
                us_info = "{}"


            req  = brunix_pb2.AgentRequest(query=query, session_id=session_id, 
                                           editor_content=ed_contxt,
                                           selected_text=sel_contxt,
                                           extra_context=ext_contxt,
                                           user_info=us_info)
            
            for resp in stub.AskAgentStream(req):                    #  AskAgentStream
                asyncio.run_coroutine_threadsafe(queue.put(resp), loop).result()
        except Exception as e:
            asyncio.run_coroutine_threadsafe(queue.put(e), loop).result()
        finally:
            asyncio.run_coroutine_threadsafe(queue.put(None), loop).result()  # sentinel

    _thread_pool.submit(_producer)

    while True:
        item = await queue.get()
        if item is None:
            break
        if isinstance(item, Exception):
            raise item
        yield item


async def _stream_chat(query: str, session_id: str, req_id: str, context = {}) -> AsyncIterator[str]:

    try:
        async for resp in _iter_stream(query, session_id, context):
            if resp.is_final:
                yield _sse(_chat_chunk("", req_id, finish="stop"))
                break
            if resp.text:
                yield _sse(_chat_chunk(resp.text, req_id))
    except Exception as e:
        logger.error(f"[stream_chat] error: {e}", exc_info=True)
        yield _sse(_chat_chunk(f"[Error: {e}]", req_id, finish="stop"))

    yield _sse_done()


async def _stream_completion(query: str, session_id: str, req_id: str) -> AsyncIterator[str]:
    try:
        async for resp in _iter_stream(query, session_id):
            if resp.is_final:
                yield _sse(_completion_chunk("", req_id, finish="stop"))
                break
            if resp.text:
                yield _sse(_completion_chunk(resp.text, req_id))
    except Exception as e:
        logger.error(f"[stream_completion] error: {e}")
        yield _sse(_completion_chunk(f"[Error: {e}]", req_id, finish="stop"))

    yield _sse_done()


def _ollama_chat_chunk(token: str, done: bool) -> str:
    return json.dumps({
        "model":      PROXY_MODEL,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "message":    {"role": "assistant", "content": token},
        "done":       done,
    }) + "\n"


def _ollama_generate_chunk(token: str, done: bool) -> str:
    return json.dumps({
        "model":      PROXY_MODEL,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "response":   token,
        "done":       done,
    }) + "\n"


async def _stream_ollama_chat(query: str, session_id: str) -> AsyncIterator[str]:
    try:
        async for resp in _iter_stream(query, session_id):
            if resp.is_final:
                yield _ollama_chat_chunk("", done=True)
                break
            if resp.text:
                yield _ollama_chat_chunk(resp.text, done=False)
    except Exception as e:
        logger.error(f"[ollama_chat] error: {e}")
        yield _ollama_chat_chunk(f"[Error: {e}]", done=True)


async def _stream_ollama_generate(query: str, session_id: str) -> AsyncIterator[str]:
    try:
        async for resp in _iter_stream(query, session_id):
            if resp.is_final:
                yield _ollama_generate_chunk("", done=True)
                break
            if resp.text:
                yield _ollama_generate_chunk(resp.text, done=False)
    except Exception as e:
        logger.error(f"[ollama_generate] error: {e}")
        yield _ollama_generate_chunk(f"[Error: {e}]", done=True)


@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{
            "id": PROXY_MODEL, "object": "model", "created": 1700000000,
            "owned_by": "brunix", "permission": [], "root": PROXY_MODEL, "parent": None,
        }],
    }


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    query      = _query_from_messages(req.messages)
    #print(req.messages)
    session_id = req.session_id or "default"
    req_id     = f"chatcmpl-{uuid.uuid4().hex}"
    
    context = {}

    try:
        context = json.loads(req.user)
    except Exception as e:
        pass


    logger.info(f"[chat] session={session_id} stream={req.stream} query='{query[:80]}'")

    if not query:
        raise HTTPException(status_code=400, detail="No user message found in messages.")

    if req.stream:

        return StreamingResponse(
            _stream_chat(query, session_id, req_id, context),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    try:
        text = await _invoke_blocking(query, session_id)
    except grpc.RpcError as e:
        raise HTTPException(status_code=502, detail=f"gRPC error: {e.details()}")

    return JSONResponse(_chat_response(text, req_id))


@app.post("/v1/completions")
async def completions(req: CompletionRequest):
    query      = req.prompt if isinstance(req.prompt, str) else " ".join(req.prompt)
    session_id = req.session_id or req.user or "default"
    req_id     = f"cmpl-{uuid.uuid4().hex}"

    logger.info(f"[completion] session={session_id} stream={req.stream} prompt='{query[:80]}'")

    if not query:
        raise HTTPException(status_code=400, detail="prompt is required.")

    if req.stream:
        return StreamingResponse(
            _stream_completion(query, session_id, req_id),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    try:
        text = await _invoke_blocking(query, session_id)
    except grpc.RpcError as e:
        raise HTTPException(status_code=502, detail=f"gRPC error: {e.details()}")

    return JSONResponse(_completion_response(text, req_id))


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


@app.get("/api/tags")
async def ollama_tags():
    return {
        "models": [{
            "name": PROXY_MODEL,
            "model":PROXY_MODEL,
            "modified_at": "2024-01-01T00:00:00Z",
            "size": 0,
            "digest":"brunix",
            "details": {
                "format":  "gguf",
                "family": "brunix",
                "parameter_size":   "unknown",
                "quantization_level": "unknown",
            },
        }]
    }


@app.post("/api/chat")
async def ollama_chat(req: OllamaChatRequest):

    query      = next((m.content for m in reversed(req.messages) if m.role == "user"), "")
    session_id = req.session_id or "default"

    logger.info(f"[ollama/chat] session={session_id} stream={req.stream} query='{query[:80]}'")

    if not query:
        raise HTTPException(status_code=400, detail="No user message found.")

    if req.stream:
        return StreamingResponse(
            _stream_ollama_chat(query, session_id),
            media_type="application/x-ndjson",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    try:
        text = await _invoke_blocking(query, session_id)
    except grpc.RpcError as e:
        raise HTTPException(status_code=502, detail=f"gRPC error: {e.details()}")

    return JSONResponse({
        "model":      PROXY_MODEL,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "message":    {"role": "assistant", "content": text},
        "done":       True,
    })


@app.post("/api/generate")
async def ollama_generate(req: OllamaGenerateRequest):

    session_id = req.session_id or "default"

    logger.info(f"[ollama/generate] session={session_id} stream={req.stream} prompt='{req.prompt[:80]}'")

    if not req.prompt:
        raise HTTPException(status_code=400, detail="prompt is required.")

    if req.stream:
        return StreamingResponse(
            _stream_ollama_generate(req.prompt, session_id),
            media_type="application/x-ndjson",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    try:
        text = await _invoke_blocking(req.prompt, session_id)
    except grpc.RpcError as e:
        raise HTTPException(status_code=502, detail=f"gRPC error: {e.details()}")

    return JSONResponse({
        "model":      PROXY_MODEL,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "response":   text,
        "done":       True,
    })