from __future__ import annotations

import json
import time
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Literal

from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.responses import PlainTextResponse, StreamingResponse
from pydantic import BaseModel, ConfigDict, Field

from .config import AppConfig, ModelConfig, load_config
from .manager import ModelManager
from .modelfile import model_config_from_modelfile
from .types import BackendError, ChatParams


class Message(BaseModel):
    model_config = ConfigDict(extra="allow")

    role: Literal["system", "user", "assistant", "tool", "function", "developer"]
    content: Any

    def payload(self) -> dict[str, Any]:
        return self.model_dump(exclude_none=True)


class ChatCompletionRequest(BaseModel):
    model_config = ConfigDict(extra="allow")

    model: str | None = None
    messages: list[Message]
    stream: bool = False
    max_tokens: int | None = Field(default=None, ge=1, le=131072)
    max_completion_tokens: int | None = Field(default=None, ge=1, le=131072)
    temperature: float | None = Field(default=None, ge=0, le=2)
    top_p: float | None = Field(default=None, ge=0, le=1)
    top_k: int | None = Field(default=None, ge=0, le=100000)
    min_p: float | None = Field(default=None, ge=0, le=1)
    repeat_penalty: float | None = Field(default=None, ge=0)
    seed: int | None = None
    stop: str | list[str] | None = None
    enable_thinking: bool | None = None
    response_format: str | dict[str, Any] | None = None
    tools: list[dict[str, Any]] = Field(default_factory=list)
    tool_choice: Any = None
    keep_alive: str | int | float | None = None
    extra_body: dict[str, Any] = Field(default_factory=dict)

    def params(self, config: AppConfig) -> ChatParams:
        del config
        default = ChatParams()
        stop = [self.stop] if isinstance(self.stop, str) else (self.stop or [])
        extra = dict(self.extra_body)
        known = {
            "model", "messages", "stream", "max_tokens", "max_completion_tokens", "temperature",
            "top_p", "top_k", "min_p", "repeat_penalty", "seed", "stop", "enable_thinking",
            "response_format", "tools", "tool_choice", "keep_alive", "extra_body",
        }
        extra.update({key: value for key, value in (self.model_extra or {}).items() if key not in known})
        return ChatParams(
            max_tokens=self.max_completion_tokens or self.max_tokens or default.max_tokens,
            temperature=default.temperature if self.temperature is None else self.temperature,
            top_p=default.top_p if self.top_p is None else self.top_p,
            top_k=default.top_k if self.top_k is None else self.top_k,
            min_p=default.min_p if self.min_p is None else self.min_p,
            repeat_penalty=default.repeat_penalty if self.repeat_penalty is None else self.repeat_penalty,
            seed=self.seed,
            stop=stop,
            enable_thinking=self.enable_thinking,
            response_format=self.response_format,
            tools=self.tools,
            tool_choice=self.tool_choice,
            extra=extra,
        )


class OllamaChatRequest(BaseModel):
    model_config = ConfigDict(extra="allow")

    model: str
    messages: list[Message]
    stream: bool = True
    format: str | dict[str, Any] | None = None
    options: dict[str, Any] = Field(default_factory=dict)
    keep_alive: str | int | float | None = "5m"
    think: bool | None = None
    tools: list[dict[str, Any]] = Field(default_factory=list)


class OllamaGenerateRequest(BaseModel):
    model_config = ConfigDict(extra="allow")

    model: str
    prompt: str = ""
    suffix: str | None = None
    system: str | None = None
    template: str | None = None
    context: list[int] | None = None
    stream: bool = True
    raw: bool = False
    format: str | dict[str, Any] | None = None
    options: dict[str, Any] = Field(default_factory=dict)
    keep_alive: str | int | float | None = "5m"
    think: bool | None = None


class OllamaEmbedRequest(BaseModel):
    model: str
    input: str | list[str]
    keep_alive: str | int | float | None = "5m"


class OpenAIEmbeddingRequest(BaseModel):
    model: str
    input: str | list[str]
    encoding_format: str | None = None


class RegisterModelRequest(BaseModel):
    id: str
    path: str | None = None
    kind: str = "auto"
    enabled: bool = True
    description: str = ""
    cloud_profile: str | None = None
    remote_model: str | None = None
    template: str | None = None
    system: str | None = None
    messages: list[dict[str, Any]] = Field(default_factory=list)
    parameters: dict[str, Any] = Field(default_factory=dict)
    options: dict[str, Any] = Field(default_factory=dict)


class CreateModelRequest(BaseModel):
    model: str
    modelfile: str
    stream: bool = False


class CopyModelRequest(BaseModel):
    source: str
    destination: str


class PullModelRequest(BaseModel):
    model: str
    path: str | None = None
    kind: str = "auto"
    stream: bool = True
    repo_id: str | None = None
    revision: str | None = None
    destination: str | None = None
    allow_patterns: list[str] | None = None
    alias: str | None = None


def _content_to_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(
            str(item.get("text", "")) if isinstance(item, dict) and item.get("type") == "text" else "[image]"
            for item in content
        )
    return str(content)


def _sse(data: dict[str, Any]) -> str:
    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"


def _iso_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _ollama_params(options: dict[str, Any], fmt: str | dict[str, Any] | None, think: bool | None, tools: list[dict[str, Any]]) -> ChatParams:
    stop = options.get("stop", [])
    if isinstance(stop, str):
        stop = [stop]
    standard = {
        "num_predict": "max_tokens",
        "temperature": "temperature",
        "top_p": "top_p",
        "top_k": "top_k",
        "min_p": "min_p",
        "repeat_penalty": "repeat_penalty",
        "seed": "seed",
    }
    values: dict[str, Any] = {key: options[source] for source, key in standard.items() if source in options}
    extra = {key: value for key, value in options.items() if key not in standard and key not in {"stop"}}
    return ChatParams(
        max_tokens=int(values.get("max_tokens", 512)),
        temperature=float(values.get("temperature", 0.7)),
        top_p=float(values.get("top_p", 0.95)),
        top_k=int(values.get("top_k", 40)),
        min_p=float(values.get("min_p", 0.05)),
        repeat_penalty=float(values.get("repeat_penalty", 1.05)),
        seed=values.get("seed"),
        stop=stop,
        enable_thinking=think,
        response_format=fmt,
        tools=tools,
        extra=extra,
    )


def _usage_from_metadata(metadata: dict[str, Any]) -> dict[str, int]:
    usage = metadata.get("usage") or {}
    if usage:
        prompt = int(usage.get("prompt_tokens", usage.get("prompt_eval_count", 0)) or 0)
        completion = int(usage.get("completion_tokens", usage.get("eval_count", 0)) or 0)
        return {
            "prompt_tokens": prompt,
            "completion_tokens": completion,
            "total_tokens": int(usage.get("total_tokens", prompt + completion)),
        }
    prompt = int(metadata.get("prompt_eval_count", metadata.get("prompt_tokens", 0)) or 0)
    completion = int(metadata.get("eval_count", metadata.get("completion_tokens", 0)) or 0)
    return {"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion}


def _ollama_usage(metadata: dict[str, Any]) -> dict[str, int]:
    usage = metadata.get("usage") or {}
    prompt = int(usage.get("prompt_eval_count", usage.get("prompt_tokens", metadata.get("prompt_eval_count", 0))) or 0)
    completion = int(usage.get("eval_count", usage.get("completion_tokens", metadata.get("eval_count", 0))) or 0)
    return {"prompt_eval_count": prompt, "eval_count": completion, "eval_duration": 0, "prompt_eval_duration": 0}


def _ollama_runtime_fields(manager: ModelManager, model_id: str) -> dict[str, int]:
    state = manager.states.get(model_id)
    if not state:
        return {}
    return {
        "total_duration": int(state.last_duration_seconds * 1_000_000_000),
        "load_duration": int(state.last_load_duration_seconds * 1_000_000_000),
    }


def create_app(config_path: str | Path | None = None) -> FastAPI:
    resolved = Path(config_path or r"E:\llm-backend\config\config.json")
    config = load_config(resolved)
    manager = ModelManager(config, resolved)

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        if config.runtime.auto_load and config.runtime.default_model:
            try:
                await manager.load(config.runtime.default_model)
            except Exception:
                pass
        yield
        await manager.shutdown()

    app = FastAPI(title="LocalPilot", version="0.2.0", lifespan=lifespan)
    app.state.config = config
    app.state.manager = manager

    async def authorize(request: Request, authorization: str | None = Header(default=None)) -> None:
        del request
        expected = config.runtime.api_key
        if expected and (authorization or "").removeprefix("Bearer ") != expected:
            raise HTTPException(status_code=401, detail="invalid API key")

    @app.get("/health")
    async def health() -> dict[str, Any]:
        return {
            "status": "ok",
            "version": "0.2.0",
            "active_model": manager.active_model_id,
            "running": manager.running_models(),
        }

    @app.get("/v1/models")
    async def models(_: None = Depends(authorize)) -> dict[str, Any]:
        data = []
        for item in manager.list_models():
            data.append({**item, "created": int(item.get("modified_at") or time.time())})
        return {"object": "list", "data": data}

    @app.get("/api/version")
    async def api_version() -> dict[str, str]:
        return {"version": "0.2.0"}

    @app.get("/api/tags")
    async def tags(_: None = Depends(authorize)) -> dict[str, Any]:
        models_data = []
        for item in manager.list_models():
            details = manager.show_model(item["id"])["details"]
            models_data.append({
                "name": item["name"],
                "model": item["model"],
                "modified_at": datetime.fromtimestamp(item["modified_at"], timezone.utc).isoformat() if item["modified_at"] else _iso_now(),
                "size": item["size"],
                "digest": item["digest"],
                "details": details,
                "capabilities": item["capabilities"],
            })
        return {"models": models_data}

    @app.get("/api/ps")
    async def ps(_: None = Depends(authorize)) -> dict[str, Any]:
        models_data = []
        for item in manager.running_models():
            state = manager.states[item["id"]]
            expires_at = None
            if state.keep_alive_seconds != float("inf"):
                expires_at = datetime.fromtimestamp(
                    state.last_used + state.keep_alive_seconds, timezone.utc
                ).isoformat()
            models_data.append({
                "name": item["name"],
                "model": item["model"],
                "size": item["size"],
                "size_vram": item.get("vram_allocated_bytes", 0),
                "expires_at": expires_at,
                "active_requests": item["active_requests"],
                "kind": item["kind"],
            })
        return {"models": models_data}

    @app.get("/metrics", response_class=PlainTextResponse)
    async def metrics(_: None = Depends(authorize)) -> str:
        records = manager.list_models()
        loaded = sum(1 for item in records if item["loaded"])
        active = sum(int(item["active_requests"]) for item in records)
        total = sum(int(item["total_requests"]) for item in records)
        lines = [
            "# TYPE localpilot_models_loaded gauge",
            f"localpilot_models_loaded {loaded}",
            "# TYPE localpilot_active_requests gauge",
            f"localpilot_active_requests {active}",
            "# TYPE localpilot_requests_total counter",
            f"localpilot_requests_total {total}",
        ]
        return "\n".join(lines) + "\n"

    @app.post("/api/show")
    async def show(body: dict[str, Any], _: None = Depends(authorize)) -> dict[str, Any]:
        model_id = body.get("model") or body.get("name")
        if not model_id:
            raise HTTPException(status_code=400, detail="missing model/name")
        try:
            return manager.show_model(model_id)
        except KeyError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc

    @app.post("/api/models/register")
    async def register(body: RegisterModelRequest, _: None = Depends(authorize)) -> dict[str, Any]:
        try:
            spec = manager.register_model(ModelConfig(**body.model_dump()))
            return {"status": "registered", "model": spec.model_dump(exclude_none=True)}
        except BackendError as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from exc

    @app.delete("/api/models/{model_id}")
    async def delete_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]:
        try:
            await manager.remove_model(model_id)
            return {"status": "deleted", "model": model_id, "files_removed": False}
        except BackendError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc

    @app.post("/api/create")
    async def create_model(body: CreateModelRequest, _: None = Depends(authorize)) -> Any:
        try:
            modelfile_text = body.modelfile
            modelfile_path = Path(modelfile_text)
            if modelfile_path.is_file():
                modelfile_text = modelfile_path.read_text(encoding="utf-8")
            spec = model_config_from_modelfile(body.model, modelfile_text)
            manager.register_model(spec)
        except (BackendError, ValueError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        if body.stream:
            async def create_stream():
                yield json.dumps({"status": "success", "model": body.model}, ensure_ascii=False) + "\n"

            return StreamingResponse(create_stream(), media_type="application/x-ndjson")
        return {"status": "success", "model": body.model}

    @app.post("/api/copy")
    async def copy_model(body: CopyModelRequest, _: None = Depends(authorize)) -> dict[str, Any]:
        try:
            spec = manager.copy_model(body.source, body.destination)
            return {"status": "success", "model": spec.id}
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    @app.post("/api/pull")
    async def pull_model(body: PullModelRequest, _: None = Depends(authorize)) -> Any:
        try:
            if body.path:
                spec = manager.register_model(ModelConfig(id=body.model, kind=body.kind, path=body.path))
            elif body.repo_id or "/" in body.model:
                repo_id = body.repo_id or body.model
                model_id = body.alias or repo_id.replace("/", "-")
                spec = await manager.pull_huggingface(
                    model_id=model_id,
                    repo_id=repo_id,
                    destination=body.destination,
                    revision=body.revision,
                    allow_patterns=body.allow_patterns,
                )
            else:
                raise HTTPException(
                    status_code=501,
                    detail="请提供 path 导入本地模型，或使用 repo_id/repo 触发 Hugging Face 下载。",
                )
        except BackendError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        if body.stream:
            async def pull_stream():
                yield json.dumps({"status": "success", "model": spec.id, "path": spec.path}, ensure_ascii=False) + "\n"

            return StreamingResponse(pull_stream(), media_type="application/x-ndjson")
        return {"status": "success", "model": spec.id, "path": spec.path}

    @app.post("/api/models/{model_id}/load")
    async def load_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]:
        try:
            backend = await manager.load(model_id)
            return backend.info()
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    @app.post("/api/models/{model_id}/unload")
    async def unload_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]:
        await manager.unload(model_id)
        return {"status": "unloaded", "model": model_id}

    @app.post("/v1/chat/completions")
    async def chat_completions(body: ChatCompletionRequest, _: None = Depends(authorize)) -> Any:
        messages = [item.payload() for item in body.messages]
        params = body.params(config)
        request_id = f"chatcmpl-localpilot-{uuid.uuid4().hex[:12]}"
        model_id = body.model or config.runtime.default_model or manager.active_model_id or "localpilot"
        try:
            if body.stream:
                async def event_stream():
                    yield _sse({
                        "id": request_id,
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "model": model_id,
                        "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
                    })
                    async for chunk in manager.chat(model_id, messages, params, stream=True, keep_alive=body.keep_alive):
                        delta = dict(chunk.delta)
                        if chunk.text and "content" not in delta:
                            delta["content"] = chunk.text
                        if delta:
                            yield _sse({
                                "id": request_id,
                                "object": "chat.completion.chunk",
                                "created": int(time.time()),
                                "model": model_id,
                                "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
                            })
                        if chunk.done:
                            yield _sse({
                                "id": request_id,
                                "object": "chat.completion.chunk",
                                "created": int(time.time()),
                                "model": model_id,
                                "choices": [{"index": 0, "delta": {}, "finish_reason": chunk.finish_reason or "stop"}],
                            })
                    yield "data: [DONE]\n\n"

                return StreamingResponse(event_stream(), media_type="text/event-stream")

            parts: list[str] = []
            finish_reason = "stop"
            metadata: dict[str, Any] = {}
            async for chunk in manager.chat(model_id, messages, params, stream=False, keep_alive=body.keep_alive):
                parts.append(chunk.text)
                finish_reason = chunk.finish_reason or finish_reason
                metadata = chunk.metadata or metadata
            message: dict[str, Any] = {"role": "assistant", "content": "".join(parts)}
            if metadata:
                original_message = ((metadata.get("choices") or [{}])[0].get("message") or {})
                for key in ("tool_calls", "function_call", "refusal"):
                    if key in original_message:
                        message[key] = original_message[key]
            return {
                "id": request_id,
                "object": "chat.completion",
                "created": int(time.time()),
                "model": model_id,
                "choices": [{
                    "index": 0,
                    "message": message,
                    "finish_reason": finish_reason,
                }],
                "usage": _usage_from_metadata(metadata),
            }
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    @app.post("/v1/completions")
    async def completions(body: dict[str, Any], _: None = Depends(authorize)) -> Any:
        prompt = body.get("prompt", "")
        if isinstance(prompt, list):
            prompt = "\n".join(str(item) for item in prompt)
        payload = ChatCompletionRequest(
            model=body.get("model"),
            messages=[Message(role="user", content=str(prompt))],
            stream=bool(body.get("stream", False)),
            max_tokens=body.get("max_tokens"),
            temperature=body.get("temperature"),
            top_p=body.get("top_p"),
            stop=body.get("stop"),
            keep_alive=body.get("keep_alive"),
        )
        return await chat_completions(payload)

    @app.post("/v1/embeddings")
    async def openai_embeddings(body: OpenAIEmbeddingRequest, _: None = Depends(authorize)) -> dict[str, Any]:
        inputs = [body.input] if isinstance(body.input, str) else body.input
        try:
            vectors = await manager.embed(body.model, inputs)
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        return {
            "object": "list",
            "data": [{"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors)],
            "model": body.model,
            "usage": {"prompt_tokens": 0, "total_tokens": 0},
        }

    @app.post("/api/embed")
    async def ollama_embed(body: OllamaEmbedRequest, _: None = Depends(authorize)) -> dict[str, Any]:
        inputs = [body.input] if isinstance(body.input, str) else body.input
        try:
            vectors = await manager.embed(body.model, inputs, body.keep_alive)
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        return {"model": body.model, "embeddings": vectors}

    @app.post("/api/embeddings")
    async def ollama_embeddings_legacy(body: OllamaEmbedRequest, _: None = Depends(authorize)) -> dict[str, Any]:
        return await ollama_embed(body)

    @app.post("/api/chat")
    async def ollama_chat(body: OllamaChatRequest, _: None = Depends(authorize)) -> Any:
        messages = [item.payload() for item in body.messages]
        params = _ollama_params(body.options, body.format, body.think, body.tools)
        try:
            if body.stream:
                async def event_stream():
                    answer: list[str] = []
                    async for chunk in manager.chat(body.model, messages, params, stream=True, keep_alive=body.keep_alive):
                        answer.append(chunk.text)
                        message = {"role": "assistant", "content": chunk.text, **chunk.delta}
                        yield json.dumps({
                            "model": body.model,
                            "created_at": _iso_now(),
                            "message": message,
                            "done": False,
                        }, ensure_ascii=False) + "\n"
                    yield json.dumps({
                        "model": body.model,
                        "created_at": _iso_now(),
                        "message": {"role": "assistant", "content": ""},
                        "done": True,
                        "done_reason": "stop",
                        **_ollama_runtime_fields(manager, body.model),
                    }, ensure_ascii=False) + "\n"

                return StreamingResponse(event_stream(), media_type="application/x-ndjson")

            parts: list[str] = []
            metadata: dict[str, Any] = {}
            async for chunk in manager.chat(body.model, messages, params, stream=False, keep_alive=body.keep_alive):
                parts.append(chunk.text)
                metadata = chunk.metadata or metadata
            return {
                "model": body.model,
                "created_at": _iso_now(),
                "message": {"role": "assistant", "content": "".join(parts)},
                "done": True,
                "done_reason": "stop",
                **_ollama_runtime_fields(manager, body.model),
                **_ollama_usage(metadata),
            }
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    @app.post("/api/generate")
    async def ollama_generate(body: OllamaGenerateRequest, _: None = Depends(authorize)) -> Any:
        del body.suffix, body.template, body.context, body.raw
        messages: list[dict[str, Any]] = []
        if body.system:
            messages.append({"role": "system", "content": body.system})
        messages.append({"role": "user", "content": body.prompt})
        params = _ollama_params(body.options, body.format, body.think, [])
        try:
            if body.stream:
                async def event_stream():
                    async for chunk in manager.chat(body.model, messages, params, stream=True, keep_alive=body.keep_alive):
                        yield json.dumps({
                            "model": body.model,
                            "created_at": _iso_now(),
                            "response": chunk.text,
                            "done": False,
                        }, ensure_ascii=False) + "\n"
                    yield json.dumps({
                        "model": body.model,
                        "created_at": _iso_now(),
                        "response": "",
                        "done": True,
                        "done_reason": "stop",
                        **_ollama_runtime_fields(manager, body.model),
                    }, ensure_ascii=False) + "\n"

                return StreamingResponse(event_stream(), media_type="application/x-ndjson")

            parts: list[str] = []
            async for chunk in manager.chat(body.model, messages, params, stream=False, keep_alive=body.keep_alive):
                parts.append(chunk.text)
            return {
                "model": body.model,
                "created_at": _iso_now(),
                "response": "".join(parts),
                "done": True,
                **_ollama_runtime_fields(manager, body.model),
            }
        except (BackendError, KeyError) as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    return app