import asyncio
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path

import httpx
from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import ValidationError
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.ai_client import active_model_name, chat_completion, strip_code_fence
from app.core.ocr import OcrUnavailableError, screenshot_to_layout_json
from app.core.config import settings
from app.core.database import get_db
from app.core.security import AuthContext, get_auth_context
from app.crud.audit import write_audit_log
from app.crud.base import CRUDBase
from app.models.audit import AiAssistantSession
from app.models.master import Master
from app.models.reporting import Dataset
from app.models.upload import UploadedFile
from app.models.workflow import Form
from app.schemas.audit import (
    AiAssistantMessageCreate,
    AiAssistantMessageResult,
    AiAssistantSessionCreate,
    AiAssistantSessionRead,
)
from app.schemas.navigation import PageBlockDef

router = APIRouter(prefix="/ai-assistant/sessions", tags=["ai-assistant"])
# Separate from `router` only because its path ("/ai-assistant/config") is a
# sibling of "/ai-assistant/sessions", not a child of it.
config_router = APIRouter(prefix="/ai-assistant", tags=["ai-assistant"])
crud = CRUDBase[AiAssistantSession, AiAssistantSessionCreate, AiAssistantSessionCreate](AiAssistantSession)

_MASTER_SCHEMA_INSTRUCTIONS = """You are drafting the field schema for a FlowFabric Master (a reusable reference data set).
Respond with ONLY a JSON object of this exact shape, no prose outside it:
{"message": "<one short sentence describing what you drafted or changed>",
 "fields": [{"name": "<field_name>", "alias": "<Display Name>", "type": "<see allowed types>", "required": true|false, "options": ["..."], "max_length": null,
             "control": null, "validation_rules": []}]}
"alias" is the human-readable display name people will see for the field instead of its technical name (e.g. name
"dist_nm" -> alias "District Name"); write it in Title Case, in the user's own wording where they gave one, and never
leave it empty. The "name" stays the short lower_snake_case key.
Allowed types: text, varchar, char, number, smallint, integer, bigint, numeric, real, double_precision,
boolean, date, time, timestamp, timestamptz, interval, uuid, json, jsonb, inet, single_select.
Prefer "text" unless a more specific type is clearly right (e.g. "integer" for a count, "date" for a due date,
"uuid" for a foreign identifier). "options" is only meaningful (and should only be non-empty) when type is
"single_select". "max_length" is only meaningful for "varchar"/"char". Keep field names short, lower_snake_case.
"control" is how the field is edited; leave it null for the type's default, or pick one that fits the type:
text: textarea, text, email, url, phone, password, rich_text, color; varchar/char: text, email, url, phone, password, color;
number types: number, currency, percentage, slider, rating; boolean: checkbox, toggle, yes_no;
single_select: dropdown, radio, autocomplete. Use email/url/phone for such values (with a varchar or text type).
"validation_rules" is a list of {"type": "<rule>", "value": <limit or null>, "message": null}; rules: min_length,
max_length, pattern (a regex), email, url, phone (no value), min_value, max_value, min_date, max_date (YYYY-MM-DD or
today/today+N/today-N). Add rules only where the user asks for them or they're clearly implied (e.g. an email field
gets {"type": "email", "value": null}, a percentage 0-100 gets min_value 0 and max_value 100)."""

_FORM_SCHEMA_INSTRUCTIONS = """You are drafting the field schema for a FlowFabric Form (a data-entry screen).
Respond with ONLY a JSON object of this exact shape, no prose outside it:
{"message": "<one short sentence describing what you drafted or changed>",
 "fields": [{"name": "<field_name>", "label": "<Human Label>", "type": "text|number|dropdown|date|checkbox|textarea|file_upload|section",
             "required": true|false, "options": ["..."], "visibility_rules": []}]}
"options" is only meaningful for type "dropdown". visibility_rules stays an empty list unless the user explicitly describes a conditional field."""

# Task Designer Studio's Custom Node runtime contract (app/crud/tasks.py's
# _run_custom_code_sync) — generated code MUST follow this exactly or the
# node will fail at run time: it's a standalone script run via subprocess.run
# in a separate OS process (not the AI's own sandbox), given the node's
# *resolved input params only* (never the whole shared context) as a JSON
# object on stdin, and it must print its result as JSON on the LAST
# non-empty line of stdout — anything else on stdout is ignored, and a
# non-zero exit or a timeout (15s) fails the node.
_TASK_CODE_PYTHON_INSTRUCTIONS = """You are writing the body of a Custom Node (Python) for FlowFabric's Task Designer Studio.
Respond with ONLY a JSON object of this exact shape, no prose outside it:
{"message": "<one short sentence describing what the script does>", "code": "<the full Python script>"}
The script runs standalone via `python script.py` in a fresh subprocess (15s timeout, no network/filesystem guarantees
beyond the OS user's own access — it is not a sandbox). It receives the node's resolved input params as a JSON object
on stdin and must print its result as JSON on the LAST line it writes to stdout — that line becomes the node's output;
everything else on stdout is ignored. A non-zero exit code fails the node. Example shape:
import sys, json
context = json.load(sys.stdin)
# ... your logic here, reading whatever keys the node's input params declare ...
print(json.dumps({"ok": True}))"""

_TASK_CODE_NODE_INSTRUCTIONS = """You are writing the body of a Custom Node (Node.js) for FlowFabric's Task Designer Studio.
Respond with ONLY a JSON object of this exact shape, no prose outside it:
{"message": "<one short sentence describing what the script does>", "code": "<the full Node.js script>"}
The script runs standalone via `node script.js` in a fresh subprocess (15s timeout, no network/filesystem guarantees
beyond the OS user's own access — it is not a sandbox). It receives the node's resolved input params as a JSON object
on stdin and must print its result as JSON on the LAST line it writes to stdout — that line becomes the node's output;
everything else on stdout is ignored. A non-zero exit code fails the node. Example shape:
let input = "";
process.stdin.on("data", (d) => (input += d));
process.stdin.on("end", () => {
  const context = JSON.parse(input);
  // ... your logic here, reading whatever keys the node's input params declare ...
  console.log(JSON.stringify({ ok: true }));
});"""


# Page Designer's screenshot-to-blocks AI Assist. Unlike the other three
# target types, this one can carry an attached image (see send_message's
# image_file_id handling) and needs to teach the model the recursive
# PageBlockDef tree shape (app/schemas/navigation.py) rather than a flat
# field list. Follows the Data Explorer's "AI Query" precedent
# (app/routers/datasets.py's ai_query_datasets/_build_catalog): the model
# is only ever given real entity *names*, never asked to invent or copy a
# UUID — "data_source_name" gets resolved against the tenant's actual
# Masters/Forms/Datasets server-side in _resolve_block_data_sources below,
# and a name with no match (or no name at all) just leaves that block's
# data_source unset for the user to wire up manually, the same as any
# blank block they'd add by hand.
_PAGE_SCHEMA_INSTRUCTIONS_TEMPLATE = """You are drafting the block layout for a FlowFabric Page (an internal \
screen built from a tree of typed blocks) that reproduces a layout the user attached a screenshot of.

You never see the screenshot's pixels. Instead you're given a JSON array of the text OCR detected in it — \
each element is {{"text": "<the detected line>", "left": <0..1>, "top": <0..1>, "width": <0..1>, "height": \
<0..1>}}, with the box position given as a FRACTION of the image's own width/height (0,0 is the top-left \
corner), roughly in reading order. Use each element's own position to infer the layout: elements with a \
similar "top" sit on the same visual row (put them in a "row" container, or as sibling fields in the \
same "form"/"section"); a taller "height" relative to nearby elements usually marks a heading; a larger \
"left" than the elements above/below it usually marks indentation (nesting inside a "section"/"panel"/\
"tabs" the way the screenshot visually groups it). OCR is imperfect — a mis-read character or two in a label \
is fine and not worth dwelling on; reproduce the STRUCTURE the positions describe, not a transcription \
exercise. Where no text was detected at all (an icon, a photo, empty whitespace) there is nothing to work \
from for that area — leave it out rather than guessing.

Respond with ONLY a JSON object of this exact shape, no prose outside it:
{{"message": "<one short sentence describing what you drafted>", "blocks": [<PageBlockDef>, ...]}}

Each PageBlockDef is an object: {{"id": "<a short unique string you invent, e.g. \\"b1\\">", "type": "<see allowed \
types below>", ...the type's own fields}}. Every "id" in the whole tree you return must be unique.

Allowed "type" values and their fields:
- "text": "text" (the string shown), "text_style" one of "heading"|"subheading"|"paragraph".
- "image": "image_url" — leave it "" unless the screenshot itself names a real URL; never invent one.
- "button": "label", "button_style" one of "primary"|"secondary"|"danger". Never include "action" — you cannot \
know which real page/form/workflow a button should open, so leave that entirely for the user to wire up afterward.
- "data_list": a table/grid of rows in the screenshot. Set "label" to a short title and "columns": [] (unknown \
yet). Only if a caption/heading right next to the table clearly names one of the DATASOURCES below, set \
"data_source_name" to that exact name (copied verbatim from the list) — otherwise omit "data_source_name" \
entirely. Leave "filters" and "row_actions" as [].
- "detail_view" / "label": same "data_source_name" rule as data_list, plus "columns": [].
- "form": a data-entry area (labeled input fields). Same "data_source_name" rule as data_list.
- "kpi_card" / "bar_chart" / "pie_chart": a number tile or chart-looking area. Same "data_source_name" rule, plus \
"color" as a hex string matching the screenshot's accent color if visible, else "#4f46e5".
- "section" / "panel": a vertically-stacked group with a heading. Set "label" to the heading text if visible, and \
"children": [<nested PageBlockDef>, ...] for whatever sits inside it, in the order it appears top to bottom.
- "row" / "column": a plain layout container with no heading of its own. "children": [<nested PageBlockDef>, ...] \
for items laid out side by side ("row") or stacked ("column").
- "tabs" / "accordion": a tabbed or collapsible-sections area. "groups": [{{"id": "<unique string>", "label": \
"<tab or section title>", "children": [<nested PageBlockDef>, ...]}}, ...], one entry per visible tab/section.

Only ever use "data_source_name" with a name copied verbatim from DATASOURCES below — never invent a name that \
isn't listed, and never construct an id yourself in any field. When nothing in DATASOURCES plausibly matches what \
a table/form/chart-like area is showing, simply omit "data_source_name" on that block; the user will connect a \
real data source themselves afterward, the same way they would for any manually-added block.

Reproduce the screenshot's visual grouping and top-to-bottom, left-to-right reading order using the container \
types ("section"/"panel"/"row"/"column"/"tabs"/"accordion") to nest things the way they're nested in the image — \
but keep the tree only as deep as the screenshot actually shows; don't invent structure that isn't visually there.

DATASOURCES (names only — copy one verbatim into "data_source_name", or omit it):
{catalog}

OCR-DETECTED TEXT ELEMENTS (from the attached screenshot, if one was attached — otherwise this section is empty \
and you're working from the user's own description alone):
{ocr_elements}
"""


async def _build_page_catalog(db: AsyncSession, tenant_id: uuid.UUID) -> tuple[str, dict[str, str]]:
    """Returns (catalog_text_for_the_prompt, {lowercased_name: "kind:id"}) —
    the same "give the model names, resolve them to real ids ourselves"
    split datasets.py's ai_query_datasets already uses. Forms are filtered
    to is_current since a Form has a composite (form_id, version) key and a
    naive "every row" query would list every historical version under the
    same name."""
    masters = list((await db.execute(select(Master).where(Master.tenant_id == tenant_id))).scalars().all())
    forms = list(
        (await db.execute(select(Form).where(Form.tenant_id == tenant_id, Form.is_current.is_(True)))).scalars().all()
    )
    datasets = list((await db.execute(select(Dataset).where(Dataset.tenant_id == tenant_id))).scalars().all())

    index: dict[str, str] = {}
    names: list[str] = []
    for m in masters:
        index[m.name.lower()] = f"master:{m.id}"
        names.append(m.name)
    for f in forms:
        index[f.name.lower()] = f"form:{f.form_id}"
        names.append(f.name)
    for d in datasets:
        index[d.name.lower()] = f"dataset:{d.id}"
        names.append(d.name)

    catalog_text = "\n".join(f"- {name}" for name in names) if names else "(no datasources, forms, or datasets exist yet)"
    return catalog_text, index


def _resolve_block_data_sources(block: dict, index: dict[str, str]) -> None:
    """Mutates `block` in place (and recurses into children/groups): turns
    a model-supplied "data_source_name" into a real "data_source" —
    "kind:id" when the name matches something in `index` (case-insensitive),
    left unset otherwise. Never lets the model's own "data_source"/id
    guesses (there shouldn't be any, per the prompt, but a model can ignore
    instructions) survive — this is the only place a data_source is set."""
    block.pop("data_source", None)
    name = block.pop("data_source_name", None)
    if isinstance(name, str) and name.strip():
        resolved = index.get(name.strip().lower())
        if resolved:
            block["data_source"] = resolved
    for child in block.get("children") or []:
        if isinstance(child, dict):
            _resolve_block_data_sources(child, index)
    for group in block.get("groups") or []:
        if isinstance(group, dict):
            for child in group.get("children") or []:
                if isinstance(child, dict):
                    _resolve_block_data_sources(child, index)


def _validate_page_blocks(raw_blocks: object) -> list[dict]:
    """A model response is untrusted input — validate each *top-level*
    block against PageBlockDef (which validates nested children/groups
    too, since they're the same recursive type) and silently drop
    whichever ones don't fit, rather than letting a malformed response
    corrupt the caller's live block tree or 500 the request."""
    if not isinstance(raw_blocks, list):
        return []
    valid: list[dict] = []
    for raw in raw_blocks:
        if not isinstance(raw, dict):
            continue
        try:
            valid.append(PageBlockDef.model_validate(raw).model_dump(mode="json", exclude_none=True))
        except ValidationError:
            continue
    return valid


def _system_prompt(target_type: str, language: str | None = None, catalog: str | None = None, ocr_elements: str | None = None) -> str:
    if target_type == "task_custom_code":
        return _TASK_CODE_NODE_INSTRUCTIONS if language == "node" else _TASK_CODE_PYTHON_INSTRUCTIONS
    if target_type == "page":
        return _PAGE_SCHEMA_INSTRUCTIONS_TEMPLATE.format(
            catalog=catalog or "(no datasources, forms, or datasets exist yet)",
            ocr_elements=ocr_elements or "(none)",
        )
    return _MASTER_SCHEMA_INSTRUCTIONS if target_type == "master" else _FORM_SCHEMA_INSTRUCTIONS


@config_router.get("/config")
async def get_ai_assistant_config() -> dict[str, str]:
    """Lets the frontend show which provider is actually active (FR-AIA-007
    extended to cover the cloud opt-in) without exposing the API key."""
    provider_label = settings.cloud_provider_label if settings.ai_provider == "cloud" else "Ollama"
    return {"provider": settings.ai_provider, "provider_label": provider_label, "model": active_model_name()}


@router.get("", response_model=list[AiAssistantSessionRead])
async def list_sessions(
    skip: int = 0,
    limit: int = 100,
    ctx: AuthContext = Depends(get_auth_context),
    db: AsyncSession = Depends(get_db),
) -> list[AiAssistantSession]:
    return await crud.list(db, ctx.tenant_id, skip=skip, limit=limit)


@router.get("/{session_id}", response_model=AiAssistantSessionRead)
async def get_session(
    session_id: uuid.UUID,
    ctx: AuthContext = Depends(get_auth_context),
    db: AsyncSession = Depends(get_db),
) -> AiAssistantSession:
    return await crud.get(db, ctx.tenant_id, session_id)


@router.post("", response_model=AiAssistantSessionRead, status_code=status.HTTP_201_CREATED)
async def create_session(
    payload: AiAssistantSessionCreate,
    ctx: AuthContext = Depends(get_auth_context),
    db: AsyncSession = Depends(get_db),
) -> AiAssistantSession:
    """FR-AIA-001/FR-AIA-006: entry point from Masters or Form Designer;
    every session records which local model answered it (FR-AIA-007)."""
    if ctx.user_id is None:
        raise HTTPException(status_code=400, detail="X-User-Id is required to start an AI Assistant session")
    session = AiAssistantSession(
        tenant_id=ctx.tenant_id,
        user_id=ctx.user_id,
        target_type=payload.target_type,
        target_id=payload.target_id,
        prompts=[],
        model_name=active_model_name(),
    )
    db.add(session)
    await db.flush()
    await db.refresh(session)
    return session


@router.post("/{session_id}/messages", response_model=AiAssistantMessageResult)
async def send_message(
    session_id: uuid.UUID,
    payload: AiAssistantMessageCreate,
    ctx: AuthContext = Depends(get_auth_context),
    db: AsyncSession = Depends(get_db),
) -> AiAssistantMessageResult:
    """
    FR-AIA-002/003: turns a natural-language instruction into a (re)drafted
    field schema, shown in the caller's live preview. FR-AIA-006: talks only
    to the self-hosted Ollama endpoint in settings.ollama_base_url — never a
    cloud AI API. FR-AIA-008: if Ollama is unreachable or returns something
    that isn't the JSON shape we asked for, surface that as ollama_error and
    leave prompts/resulting_schema exactly as they were — don't discard
    fields the user already accepted in a prior turn.
    """
    session = await crud.get(db, ctx.tenant_id, session_id)

    # "page" sessions only — an attached screenshot (first message of the
    # session) is read off disk and run through OCR (app/core/ocr.py) here:
    # the AI Assistant's own model is text-only (see Settings.ai_provider),
    # so the pixels themselves are never sent to it — only the JSON list of
    # text elements OCR detected, folded into the system prompt below. The
    # image bytes (and the OCR JSON itself) are never persisted into
    # session.prompts (see stored_text below) — read and used only for this
    # one outbound request — so a session's JSONB row stays small no matter
    # how many regenerate turns it accumulates.
    image_filename: str | None = None
    ocr_element_count = 0
    ocr_elements_json: str | None = None
    if payload.image_file_id is not None:
        upload = await CRUDBase(UploadedFile).get(db, ctx.tenant_id, payload.image_file_id)
        if not upload.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail="Attached file is not an image")
        upload_path = Path(upload.storage_path)
        if not upload_path.is_file():
            raise HTTPException(status_code=404, detail="Attached image is missing from storage")
        raw_bytes = await asyncio.to_thread(upload_path.read_bytes)
        image_filename = upload.filename
        try:
            layout = await asyncio.to_thread(screenshot_to_layout_json, raw_bytes)
        except OcrUnavailableError as exc:
            return AiAssistantMessageResult(
                session=AiAssistantSessionRead.model_validate(session),
                reply="",
                ollama_error=str(exc),
            )
        ocr_element_count = len(layout["elements"])
        ocr_elements_json = json.dumps(layout["elements"], separators=(",", ":"))

    catalog_text: str | None = None
    catalog_index: dict[str, str] = {}
    if session.target_type == "page":
        catalog_text, catalog_index = await _build_page_catalog(db, ctx.tenant_id)

    stored_text = (
        f"{payload.message}\n[attached screenshot: {image_filename} — {ocr_element_count} text region(s) read via OCR]"
        if image_filename
        else payload.message
    )
    user_message = {"role": "user", "content": stored_text, "at": datetime.now(timezone.utc).isoformat()}
    messages = [
        {
            "role": "system",
            "content": _system_prompt(session.target_type, payload.language, catalog_text, ocr_elements_json),
        },
        *[{"role": m["role"], "content": m["content"]} for m in session.prompts],
        {"role": "user", "content": payload.message},
    ]

    try:
        content = await chat_completion(messages)
        parsed = json.loads(strip_code_fence(content))
    except (httpx.HTTPError, KeyError, json.JSONDecodeError) as exc:
        return AiAssistantMessageResult(
            session=AiAssistantSessionRead.model_validate(session),
            reply="",
            ollama_error=f"AI Assistant is unavailable ({exc}); continue in the manual builder — nothing already drafted was lost.",
        )

    assistant_message = {
        "role": "assistant",
        "content": parsed.get("message", ""),
        "at": datetime.now(timezone.utc).isoformat(),
    }
    session.prompts = [*session.prompts, user_message, assistant_message]
    if session.target_type == "task_custom_code":
        session.resulting_schema = {"code": parsed.get("code", ""), "language": payload.language or "python"}
        audit_details = {"target_type": session.target_type, "code_length": len(parsed.get("code", ""))}
    elif session.target_type == "page":
        raw_blocks = parsed.get("blocks", [])
        if isinstance(raw_blocks, list):
            for raw_block in raw_blocks:
                if isinstance(raw_block, dict):
                    _resolve_block_data_sources(raw_block, catalog_index)
        blocks = _validate_page_blocks(raw_blocks)
        session.resulting_schema = {"blocks": blocks}
        audit_details = {"target_type": session.target_type, "block_count": len(blocks)}
    else:
        session.resulting_schema = {"fields": parsed.get("fields", [])}
        audit_details = {"target_type": session.target_type, "field_count": len(parsed.get("fields", []))}
    await db.flush()
    await db.refresh(session)

    await write_audit_log(
        db, ctx.tenant_id, ctx.user_id, "ai_assistant.message", "ai_assistant_session", session_id, audit_details,
    )

    return AiAssistantMessageResult(
        session=AiAssistantSessionRead.model_validate(session),
        reply=parsed.get("message", ""),
    )
