import os

import httpx
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel


class PromptRequest(BaseModel):
    prompt: str
    stream: bool = False


app = FastAPI()
# Ollama lives on SOLARIA (GPU node); resolved via Tailscale MagicDNS on the
# host, which Docker's embedded DNS forwards to. Overridable per node via env.
OLLAMA_URL = os.environ.get('OLLAMA_URL', 'http://solaria:11434/api/generate')
CHAT_MODEL = os.environ.get('CHAT_MODEL', 'deepcoder:14b')
CODE_MODEL = os.environ.get('CODE_MODEL', 'deepseek-coder:latest')
NATURAL_LANGUAGE_STARTERS = {
    'write',
    'create',
    'generate',
    'make',
    'say',
    'show',
    'give',
    'explain',
    'tell',
    'build',
}


@app.get('/')
async def root() -> dict[str, str]:
    return {'status': 'gateway ok'}


async def generate(route: str, model: str, request: PromptRequest) -> dict[str, object]:
    payload = {
        'model': model,
        'prompt': request.prompt,
        'stream': False,
    }
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            ollama_response = await client.post(OLLAMA_URL, json=payload)
            ollama_response.raise_for_status()
    except httpx.HTTPStatusError as exc:
        raise HTTPException(status_code=502, detail=f'Ollama error: {exc.response.text}') from exc
    except httpx.HTTPError as exc:
        raise HTTPException(status_code=502, detail=f'Ollama unreachable: {exc}') from exc

    raw = ollama_response.json()
    return {
        'route': route,
        'model': model,
        'response': raw.get('response', ''),
        'raw': raw,
    }


def looks_like_shell_command(prompt: str) -> bool:
    stripped = prompt.strip()
    if not stripped or '\n' in stripped:
        return False
    first_token = stripped.split()[0].lower()
    if first_token in NATURAL_LANGUAGE_STARTERS:
        return False
    return True


@app.post('/api/chat')
async def chat(request: PromptRequest) -> dict[str, object]:
    return await generate('chat', CHAT_MODEL, request)


@app.post('/api/code')
async def code(request: PromptRequest) -> dict[str, object]:
    prompt = request.prompt
    code_prompt = f"""
You are a coding and command assistant.

Return ONLY the final answer.
Do NOT explain.
Do NOT use markdown.
Do NOT mention that you are an AI.
Do NOT add prose.

If the task is a shell command request, return only the shell command.
If the task is a programming request, return only the code.

TASK:
{prompt}
""".strip()
    code_request = PromptRequest(prompt=code_prompt, stream=request.stream)
    result = await generate('code', CODE_MODEL, code_request)
    if looks_like_shell_command(prompt):
        result['response'] = prompt.strip()
    return result