import os
import time
import logging
import asyncio
from datetime import datetime
from typing import AsyncIterator
import ollama
from dotenv import load_dotenv

from prompts import get_vlm_ocr_prompt

load_dotenv()

OLLAMA_MODEL = os.getenv('OLLAMA_MODEL', 'gpt-oss:20b')
PRO_OLLAMA_MODEL = os.getenv('PRO_OLLAMA_MODEL', OLLAMA_MODEL)
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
VLM_MODEL = os.getenv('VLM_MODEL', 'qwen3-vl:30b')

# Timeouts in seconds (10 minutes for large model loading)
COMPLETION_TIMEOUT = 600
OCR_TIMEOUT = 600

client = ollama.AsyncClient(host=OLLAMA_HOST)
logger = logging.getLogger("llm")


def _extract_message(response) -> tuple[str, str]:
    content = ""
    thinking = ""

    if hasattr(response, 'message') and response.message:
        content = response.message.content or ""
        thinking = getattr(response.message, 'thinking', '') or ""
    elif isinstance(response, dict) and 'message' in response:
        msg = response.get('message', {})
        content = msg.get('content', '') or ""
        thinking = msg.get('thinking', '') or ""

    # fallback for generate
    if not content:
        if hasattr(response, 'response'):
            content = getattr(response, 'response', '') or ""
        elif isinstance(response, dict) and 'response' in response:
            content = response.get('response', '') or ""

    return content, thinking


def _build_prompt(prompt: str, system_prompt: str | None = None) -> str:
    if system_prompt and system_prompt.strip():
        return f"{system_prompt}\n\n{prompt}"
    return prompt


def _resolve_model_name(model: str | None = None, *, use_pro_model: bool = False) -> str:
    candidate = (model or '').strip()
    if candidate:
        return candidate
    return PRO_OLLAMA_MODEL if use_pro_model else OLLAMA_MODEL


def _build_generate_kwargs(
    prompt: str,
    *,
    system_prompt: str | None = None,
    temperature: float = 0.7,
    thinking: str | None = None,
    model: str | None = None,
    use_pro_model: bool = False,
    stream: bool = False,
) -> dict:
    kwargs = {
        "model": _resolve_model_name(model, use_pro_model=use_pro_model),
        "prompt": _build_prompt(prompt, system_prompt),
        "stream": stream,
        "raw": True,
        "options": {
            'temperature': temperature,
            'repeat_penalty': 1.1,
        },
    }
    if thinking:
        kwargs["think"] = thinking
    return kwargs


def _extract_stream_text(chunk) -> str:
    content, _ = _extract_message(chunk)
    if content:
        return content

    if isinstance(chunk, dict):
        return chunk.get('response', '') or ''

    if hasattr(chunk, 'response'):
        return getattr(chunk, 'response', '') or ''

    return ''


async def call_ollama(
    prompt: str,
    *,
    system_prompt: str | None = None,
    tag: str = "default",
    temperature: float = 0.7,
    thinking: str | None = None,
    model: str | None = None,
    use_pro_model: bool = False,
) -> dict:
    """
    调用 Ollama API 并返回 content 和 thinking。
    """
    start = time.perf_counter()
    start_dt = datetime.now()
    model_name = _resolve_model_name(model, use_pro_model=use_pro_model)
    logger.info(
        "[LLM][%s] request model=%s host=%s prompt_chars=%d system_chars=%d temp=%.2f thinking=%s",
        tag,
        model_name,
        OLLAMA_HOST,
        len(prompt),
        len(system_prompt or ""),
        temperature,
        thinking,
    )

    try:
        kwargs = _build_generate_kwargs(
            prompt,
            system_prompt=system_prompt,
            temperature=temperature,
            thinking=thinking,
            model=model,
            use_pro_model=use_pro_model,
            stream=False,
        )

        response = await asyncio.wait_for(client.generate(**kwargs), timeout=COMPLETION_TIMEOUT)
    except asyncio.CancelledError:
        elapsed_ms = (time.perf_counter() - start) * 1000
        end_dt = datetime.now()
        logger.info(
            "[LLM][%s] call_time [%s --> %s]",
            tag,
            start_dt.strftime("%H:%M:%S"),
            end_dt.strftime("%H:%M:%S"),
        )
        logger.warning("[LLM][%s] request cancelled after %.1fms", tag, elapsed_ms)
        raise
    except Exception:
        elapsed_ms = (time.perf_counter() - start) * 1000
        end_dt = datetime.now()
        logger.info(
            "[LLM][%s] call_time [%s --> %s]",
            tag,
            start_dt.strftime("%H:%M:%S"),
            end_dt.strftime("%H:%M:%S"),
        )
        logger.exception("[LLM][%s] request failed after %.1fms", tag, elapsed_ms)
        raise

    content, thinking = _extract_message(response)
    elapsed_ms = (time.perf_counter() - start) * 1000
    end_dt = datetime.now()
    logger.info(
        "[LLM][%s] call_time [%s --> %s]",
        tag,
        start_dt.strftime("%H:%M:%S"),
        end_dt.strftime("%H:%M:%S"),
    )
    logger.info(
        "[LLM][%s] response in %.1fms response_type=%s content_chars=%d thinking_chars=%d",
        tag,
        elapsed_ms,
        type(response).__name__,
        len(content),
        len(thinking),
    )

    if not content.strip():
        logger.warning("[LLM][%s] empty content returned by model", tag)

    return {"content": content, "think": thinking}


async def stream_ollama(
    prompt: str,
    *,
    system_prompt: str | None = None,
    tag: str = "default-stream",
    temperature: float = 0.7,
    thinking: str | None = None,
    model: str | None = None,
    use_pro_model: bool = False,
) -> AsyncIterator[str]:
    start = time.perf_counter()
    start_dt = datetime.now()
    model_name = _resolve_model_name(model, use_pro_model=use_pro_model)
    yielded_chars = 0

    logger.info(
        "[LLM][%s] stream request model=%s host=%s prompt_chars=%d system_chars=%d temp=%.2f thinking=%s",
        tag,
        model_name,
        OLLAMA_HOST,
        len(prompt),
        len(system_prompt or ""),
        temperature,
        thinking,
    )

    try:
        kwargs = _build_generate_kwargs(
            prompt,
            system_prompt=system_prompt,
            temperature=temperature,
            thinking=thinking,
            model=model,
            use_pro_model=use_pro_model,
            stream=True,
        )
        stream = await client.generate(**kwargs)
        iterator = stream.__aiter__()
        deadline = time.perf_counter() + COMPLETION_TIMEOUT

        while True:
            remaining = deadline - time.perf_counter()
            if remaining <= 0:
                raise TimeoutError("LLM stream timed out")

            try:
                chunk = await asyncio.wait_for(iterator.__anext__(), timeout=remaining)
            except StopAsyncIteration:
                break

            text = _extract_stream_text(chunk)
            if not text:
                continue

            yielded_chars += len(text)
            yield text
    except asyncio.CancelledError:
        elapsed_ms = (time.perf_counter() - start) * 1000
        end_dt = datetime.now()
        logger.info(
            "[LLM][%s] stream_time [%s --> %s]",
            tag,
            start_dt.strftime("%H:%M:%S"),
            end_dt.strftime("%H:%M:%S"),
        )
        logger.warning("[LLM][%s] stream cancelled after %.1fms", tag, elapsed_ms)
        raise
    except Exception:
        elapsed_ms = (time.perf_counter() - start) * 1000
        end_dt = datetime.now()
        logger.info(
            "[LLM][%s] stream_time [%s --> %s]",
            tag,
            start_dt.strftime("%H:%M:%S"),
            end_dt.strftime("%H:%M:%S"),
        )
        logger.exception("[LLM][%s] stream failed after %.1fms", tag, elapsed_ms)
        raise

    elapsed_ms = (time.perf_counter() - start) * 1000
    end_dt = datetime.now()
    logger.info(
        "[LLM][%s] stream_time [%s --> %s]",
        tag,
        start_dt.strftime("%H:%M:%S"),
        end_dt.strftime("%H:%M:%S"),
    )
    logger.info(
        "[LLM][%s] stream finished in %.1fms yielded_chars=%d",
        tag,
        elapsed_ms,
        yielded_chars,
    )

async def call_vlm_ocr(image_bytes: bytes, language: str = 'auto') -> str:
    start = time.perf_counter()
    start_dt = datetime.now()
    logger.info(
        "[VLM][ocr] request model=%s host=%s image_bytes=%d language=%s",
        VLM_MODEL,
        OLLAMA_HOST,
        len(image_bytes),
        language,
    )

    try:
        response = await asyncio.wait_for(
            client.chat(
                model=VLM_MODEL,
                messages=[{
                    'role': 'user',
                    'content': get_vlm_ocr_prompt(),
                    'images': [image_bytes]
                }],
                stream=False,
                options={'temperature': 0.3}
            ),
            timeout=OCR_TIMEOUT
        )
    except Exception:
        elapsed_ms = (time.perf_counter() - start) * 1000
        end_dt = datetime.now()
        logger.info(
            "[VLM][ocr] call_time [%s --> %s]",
            start_dt.strftime("%H:%M:%S"),
            end_dt.strftime("%H:%M:%S"),
        )
        logger.exception("[VLM][ocr] request failed after %.1fms", elapsed_ms)
        raise

    content, thinking = _extract_message(response)
    elapsed_ms = (time.perf_counter() - start) * 1000
    end_dt = datetime.now()
    logger.info(
        "[VLM][ocr] call_time [%s --> %s]",
        start_dt.strftime("%H:%M:%S"),
        end_dt.strftime("%H:%M:%S"),
    )
    logger.info(
        "[VLM][ocr] response in %.1fms response_type=%s content_chars=%d thinking_chars=%d",
        elapsed_ms,
        type(response).__name__,
        len(content),
        len(thinking),
    )

    if not content.strip():
        logger.warning("[VLM][ocr] empty content returned by model")

    return content