import asyncio
import base64
import json
import logging
import os
import re
import shutil
import subprocess
import tempfile
import uuid
from typing import Optional

from fastapi import FastAPI, HTTPException, Request, Security
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import APIKeyHeader
from pydantic import BaseModel

from geoip import get_ip_location_text
from llm import call_ollama, call_vlm_ocr, stream_ollama
from models import UserPreferences
from prompt import build_completion_prompts, prepare_prompt_context
import markitdown

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s - %(message)s",
)
logger = logging.getLogger("api")

_markitdown_instance = None


def _get_markitdown():  # pragma: no cover
    global _markitdown_instance
    if _markitdown_instance is None:
        _markitdown_instance = markitdown.MarkItDown()
    return _markitdown_instance

app = FastAPI()

# Startup event disabled — TTS model loads lazily on first request
# to avoid blocking startup and OOM crashes.
ACTIVE_COMPLETIONS: dict[str, asyncio.Task] = {}
ACTIVE_COMPLETIONS_LOCK = asyncio.Lock()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*", "X-API-Key", "X-Client-IP", "X-Request-Id"],
)

API_KEY = os.getenv("API_KEY", "your-secret-key-here")
api_key_header = APIKeyHeader(name="X-API-Key")


async def get_api_key(api_key: str = Security(api_key_header)):  # pragma: no cover
    if api_key != API_KEY:
        raise HTTPException(
            status_code=403,
            detail="Could not validate credentials",
        )
    return api_key


class CompletionRequest(BaseModel):
    prefix: str
    suffix: str
    languageId: str = "markdown"
    model_thinking: str = "low"
    privacy_mode: bool = False
    user_preferences: Optional[UserPreferences] = None
    model: Optional[str] = None
    temperature: float = 0.7


class CancelCompletionRequest(BaseModel):
    request_id: str
    reason: str = "abort"


class OCRRequest(BaseModel):
    image: str
    filename: str = "image.jpg"
    language: str = "auto"


class ConvertRequest(BaseModel):
    file: str
    filename: str = "document.pdf"


ALLOWED_CONVERT_EXTENSIONS = {".txt", ".docx", ".pptx", ".pdf"}
IMAGE_MARKDOWN_RE = re.compile(r"!\[[^\]]*]\([^)]+\)")
IMAGE_HTML_RE = re.compile(r"<img\b[^>]*>", re.IGNORECASE)


def _convert_docx_to_pdf(input_path: str, output_path: str) -> None:  # pragma: no cover
    node_executable = shutil.which("node")
    if not node_executable:
        raise RuntimeError("未找到 Node.js，无法转换 DOCX 为 PDF")

    bridge_path = os.path.join(os.path.dirname(__file__), "docx2pdf_bridge.cjs")
    if not os.path.exists(bridge_path):
        raise RuntimeError("缺少 DOCX 转 PDF 桥接脚本")

    result = subprocess.run(
        [node_executable, bridge_path, input_path, output_path],
        cwd=os.path.dirname(os.path.dirname(__file__)),
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        error_text = (result.stderr or result.stdout or "DOCX 转 PDF 失败").strip()
        raise RuntimeError(error_text)


def _preview(text: str, limit: int = 80) -> str:
    value = (text or "").replace("\n", "\\n")
    if len(value) <= limit:
        return value
    return value[:limit] + "..."


def _sanitize_converted_markdown(text: str) -> str:
    value = (text or "").replace("\r\n", "\n").replace("\r", "\n")
    value = IMAGE_MARKDOWN_RE.sub("", value)
    value = IMAGE_HTML_RE.sub("", value)
    value = re.sub(r"\n{3,}", "\n\n", value)
    return value.strip()


def get_client_ip(request: Request) -> str:
    if request.client:
        return request.headers.get("X-Client-IP") or request.client.host
    return request.headers.get("X-Client-IP") or "unknown"


def _clamp_temperature(value: float, default: float = 0.7) -> float:
    try:
        numeric = float(value)
    except (TypeError, ValueError):
        return default
    return max(0.0, min(numeric, 1.2))


@app.post("/v1/completions")
async def create_completion(request: Request, req: CompletionRequest, api_key: str = Security(get_api_key)):
    request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
    request_tag = request_id[:8]
    inference_task: Optional[asyncio.Task] = None

    client_ip = "hidden"
    location = ""

    if not req.privacy_mode:  # pragma: no cover
        client_ip = get_client_ip(request)
        location = get_ip_location_text(client_ip)
        if location:
            logger.info("[%s] client_location=%s", request_tag, location)

    try:
        logger.info(
            "[%s] /v1/completions request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d lang=%s thinking=%s privacy=%s",
            request_tag,
            request_id,
            client_ip,
            len(req.prefix or ""),
            len(req.suffix or ""),
            req.languageId,
            req.model_thinking,
            req.privacy_mode,
        )

        llm_prefix, llm_suffix = prepare_prompt_context(req.prefix or "", req.suffix or "")
        logger.info("[%s] llm_input_prefix=%r", request_tag, llm_prefix)
        logger.info("[%s] llm_input_suffix=%r", request_tag, llm_suffix)

        system_prompt, user_prompt, prefill = build_completion_prompts(
            req.prefix,
            req.suffix,
            req.languageId,
            location=location,
            thinking_level=req.model_thinking,
            preferences=req.user_preferences,
        )

        inference_task = asyncio.create_task(
            call_ollama(
                user_prompt,
                system_prompt=system_prompt,
                tag=f"{request_tag}-primary",
                temperature=_clamp_temperature(req.temperature, 0.7),
                thinking=req.model_thinking if req.model_thinking != "none" else None,
                model=req.model,
                prefill=prefill or None,
            )
        )

        existing = ACTIVE_COMPLETIONS.get(request_id)
        if existing and not existing.done():
            existing.cancel()
        ACTIVE_COMPLETIONS[request_id] = inference_task

        result = await inference_task
        content = result["content"] or ""
        if not content.strip():
            logger.warning("[%s] primary returned empty content, returning empty result", request_tag)
        logger.info(
            "[%s] completion resolved source=primary request_id=%s content_chars=%d content_preview='%s'",
            request_tag,
            request_id,
            len(content),
            _preview(content, 120),
        )

        return JSONResponse(content={"content": content, "request_id": request_id})
    except asyncio.CancelledError:
        logger.info("[%s] /v1/completions cancelled request_id=%s", request_tag, request_id)
        return JSONResponse(content={"cancelled": True, "request_id": request_id}, status_code=499)
    except Exception as e:
        logger.exception("[%s] /v1/completions failed request_id=%s: %s", request_tag, request_id, e)
        return JSONResponse(content={"error": str(e)}, status_code=500)
    finally:
        active = ACTIVE_COMPLETIONS.get(request_id)
        if active is not None and active is inference_task:
            ACTIVE_COMPLETIONS.pop(request_id, None)


@app.post("/v1/pro/completions/stream")
async def create_pro_completion_stream(request: Request, req: CompletionRequest, api_key: str = Security(get_api_key)):
    request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
    request_tag = request_id[:8]
    queue: asyncio.Queue[Optional[tuple[str, str]]] = asyncio.Queue()

    client_ip = "hidden"
    location = ""

    if not req.privacy_mode:  # pragma: no cover
        client_ip = get_client_ip(request)
        location = get_ip_location_text(client_ip)
        if location:
            logger.info("[%s] client_location=%s", request_tag, location)

    logger.info(
        "[%s] /v1/pro/completions/stream request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d lang=%s thinking=%s privacy=%s model=%s temp=%.2f",
        request_tag,
        request_id,
        client_ip,
        len(req.prefix or ""),
        len(req.suffix or ""),
        req.languageId,
        req.model_thinking,
        req.privacy_mode,
        req.model or "",
        _clamp_temperature(req.temperature, 0.7),
    )

    llm_prefix, llm_suffix = prepare_prompt_context(req.prefix or "", req.suffix or "")
    logger.info("[%s] pro_llm_input_prefix=%r", request_tag, llm_prefix)
    logger.info("[%s] pro_llm_input_suffix=%r", request_tag, llm_suffix)

    system_prompt, user_prompt, prefill = build_completion_prompts(
        req.prefix,
        req.suffix,
        req.languageId,
        location=location,
        thinking_level=req.model_thinking,
        preferences=req.user_preferences,
    )

    async def producer() -> None:
        chunks: list[str] = []
        try:
            async for delta in stream_ollama(
                user_prompt,
                system_prompt=system_prompt,
                tag=f"{request_tag}-pro",
                temperature=_clamp_temperature(req.temperature, 0.7),
                thinking=req.model_thinking if req.model_thinking != "none" else None,
                model=req.model,
                use_pro_model=True,
                prefill=prefill or None,
            ):
                chunks.append(delta)
                await queue.put(("chunk", json.dumps({"delta": delta}, ensure_ascii=False)))

            content = "".join(chunks)
            logger.info(
                "[%s] pro stream resolved request_id=%s content_chars=%d content_preview='%s'",
                request_tag,
                request_id,
                len(content),
                _preview(content, 120),
            )
            await queue.put((
                "done",
                json.dumps({"content": content, "request_id": request_id}, ensure_ascii=False),
            ))
        except asyncio.CancelledError:
            logger.info("[%s] /v1/pro/completions/stream cancelled request_id=%s", request_tag, request_id)
            await queue.put((
                "cancelled",
                json.dumps({"cancelled": True, "request_id": request_id}, ensure_ascii=False),
            ))
            raise
        except Exception as e:
            logger.exception("[%s] /v1/pro/completions/stream failed request_id=%s: %s", request_tag, request_id, e)
            await queue.put((
                "error",
                json.dumps({"error": str(e), "request_id": request_id}, ensure_ascii=False),
            ))
        finally:
            await queue.put(None)

    producer_task = asyncio.create_task(producer())
    existing = ACTIVE_COMPLETIONS.get(request_id)
    if existing and not existing.done():
        existing.cancel()
    ACTIVE_COMPLETIONS[request_id] = producer_task

    async def event_stream():
        try:
            while True:
                item = await queue.get()
                if item is None:
                    break

                event_name, data = item
                yield f"event: {event_name}\ndata: {data}\n\n"
        except asyncio.CancelledError:
            producer_task.cancel()
            raise
        finally:
            active = ACTIVE_COMPLETIONS.get(request_id)
            if active is producer_task:
                ACTIVE_COMPLETIONS.pop(request_id, None)

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


@app.post("/v1/completions/cancel")
async def cancel_completion(req: CancelCompletionRequest, api_key: str = Security(get_api_key)):
    request_tag = str(uuid.uuid4())[:8]
    request_id = req.request_id or ""

    async with ACTIVE_COMPLETIONS_LOCK:
        task = ACTIVE_COMPLETIONS.get(request_id)
        if task is None:
            logger.info(
                "[%s] /v1/completions/cancel request_id=%s status=not_found reason=%s",
                request_tag,
                request_id,
                req.reason,
            )
            return {"cancelled": False, "status": "not_found"}

        if task.done():
            logger.info(
                "[%s] /v1/completions/cancel request_id=%s status=already_done reason=%s",
                request_tag,
                request_id,
                req.reason,
            )
            return {"cancelled": False, "status": "already_done"}

        task.cancel()

    logger.info(
        "[%s] /v1/completions/cancel request_id=%s status=ok reason=%s",
        request_tag,
        request_id,
        req.reason,
    )
    return {"cancelled": True, "status": "ok"}


@app.post("/v1/ocr")
async def ocr_image(request: OCRRequest, api_key: str = Security(get_api_key)):
    request_id = str(uuid.uuid4())[:8]
    try:
        logger.info(
            "[%s] /v1/ocr filename=%s language=%s image_base64_chars=%d",
            request_id,
            request.filename,
            request.language,
            len(request.image or ""),
        )
        image_bytes = base64.b64decode(request.image)
        logger.info("[%s] /v1/ocr decoded image_bytes=%d", request_id, len(image_bytes))
        result = await call_vlm_ocr(image_bytes, request.language)  
        logger.info(
            "[%s] /v1/ocr success text_chars=%d text_preview='%s'",
            request_id,
            len(result or ""),
            _preview(result or "", 120),
        )
        return {"text": result, "filename": request.filename}
    except Exception as e:
        logger.exception("[%s] /v1/ocr failed: %s", request_id, e)
        return JSONResponse(content={"error": str(e)}, status_code=500)


@app.post("/v1/convert")
async def convert_to_markdown(request: ConvertRequest, api_key: str = Security(get_api_key)):
    """Convert file to markdown"""
    request_id = str(uuid.uuid4())[:8]

    try:
        logger.info(
            "[%s] /v1/convert filename=%s file_base64_chars=%d",
            request_id,
            request.filename,
            len(request.file or ""),
        )

        # Decode base64
        file_bytes = base64.b64decode(request.file)
        logger.info("[%s] /v1/convert decoded file_bytes=%d", request_id, len(file_bytes))

        # Get file extension
        ext = os.path.splitext(request.filename)[1].lower()

        if ext not in ALLOWED_CONVERT_EXTENSIONS:
            raise ValueError("仅支持 txt、docx、pptx、pdf 格式")

        if ext == ".txt":
            markdown_text = _sanitize_converted_markdown(file_bytes.decode("utf-8", errors="ignore"))
            return {
                "markdown": markdown_text,
                "filename": request.filename
            }

        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
            tmp.write(file_bytes)
            tmp_path = tmp.name

        try:
            # Convert using MarkItDown
            md = _get_markitdown()
            result = await asyncio.to_thread(md.convert, tmp_path)
            markdown_text = _sanitize_converted_markdown(result.text_content)

            logger.info(
                "[%s] /v1/convert success text_chars=%d text_preview='%s'",
                request_id,
                len(markdown_text or ""),
                _preview(markdown_text, 120),
            )

            return {
                "markdown": markdown_text,
                "filename": request.filename
            }
        finally:
            # Clean up temporary file
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    except Exception as e:
        logger.exception("[%s] /v1/convert failed: %s", request_id, e)
        return JSONResponse(content={"error": str(e)}, status_code=500)

# TTS and ASR routes (lazy loaded to avoid heavy import on startup)
def _register_tts_asr_routes():
    try:
        from tts_asr import register_tts_asr_routes
    except ModuleNotFoundError as exc:
        logger.warning("Skipping TTS/ASR route registration because a dependency is missing: %s", exc)
        return
    except Exception as exc:
        logger.warning("Skipping TTS/ASR route registration because import failed: %s", exc)
        return

    try:
        register_tts_asr_routes(app)
    except Exception as exc:
        logger.warning("Failed to register TTS/ASR routes: %s", exc)

_register_tts_asr_routes()

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8001)