"""Agent tools - callable operations for the agentic research engine."""
from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any
import httpx

from app.config import get_settings
from app.db.database import db


async def read_document(doc_id: str) -> dict[str, Any]:
    """Read complete document record."""
    doc = await db.get_document(doc_id)
    if not doc:
        return {"error": "Document not found"}
    return {"document": doc}


async def read_chunks(doc_id: str, page_num: int | None = None, limit: int = 50) -> str:
    """Read text chunks for a document."""
    chunks = await db.get_doc_chunks(doc_id)
    if page_num is not None:
        chunks = [c for c in chunks if c.get("page_num") == page_num]
    chunks = chunks[:limit]
    parts = []
    for c in chunks:
        page_part = f"[p{c['page_num']}]" if c.get("page_num") else ""
        parts.append(f"{page_part} {c['content'][:500]}")
    return "\n\n".join(parts)


async def get_page_text(doc_id: str, page_num: int) -> str:
    """Extract text for a specific page, including polygon metadata."""
    chunks = await db.get_doc_chunks(doc_id)
    page_chunks = [c for c in chunks if c.get("page_num") == page_num]
    if not page_chunks:
        return f"No content available for page {page_num}"
    
    # Build page view with polygon info
    text_parts = []
    for c in page_chunks:
        polygon = c.get("polygon")
        poly_info = ""
        if polygon and isinstance(polygon, dict):
            bbox = polygon.get("bbox", [])
            if bbox:
                poly_info = f" [bbox:{bbox[0]:.0f},{bbox[1]:.0f},{bbox[2]:.0f},{bbox[3]:.0f}]"
        text_parts.append(f"[{c['block_index']}{poly_info}] {c['content'][:300]}")
    
    return f"--- Page {page_num} ---\n" + "\n".join(text_parts)


async def vector_search(query: str, doc_id: str | None = None, limit: int = 20) -> str:
    """Search document chunks by semantic similarity."""
    settings = get_settings()
    async with httpx.AsyncClient(timeout=30) as client:
        query_vec = []
        # Try to get embedding via Ollama
        try:
            resp = await client.post(
                f"{settings.ollama_url}/api/embeddings",
                json={"model": "nomic-embed-text", "prompt": query},
            )
            query_vec = resp.json().get("embedding", [])
        except Exception:
            # Fall back to local embedding
            from app.core import get_embedding
            query_vec = get_embedding(query)
    
    if not query_vec:
        return "No embeddings available. Using text search fallback."
    
    results = await db.vector_search(query_vec, doc_id=doc_id, limit=limit)
    parts = []
    for r in results:
        score = r.get("similarity", 0)
        page = r.get("page_num", "?")
        parts.append(f"[{score:.2f}] (p{page}) {r['content'][:200]}")
    return "\n".join(parts) if parts else "No similar content found."


async def text_search(query: str, doc_id: str | None = None, limit: int = 20) -> str:
    """Search chunks by keyword match."""
    results = await db.search_chunks_text(query, doc_id=doc_id, limit=limit)
    parts = []
    for r in results:
        rank = r.get("rank", 0)
        page = r.get("page_num", "?")
        parts.append(f"[{rank:.2f}] (p{page}) {r['content'][:200]}")
    return "\n".join(parts) if parts else "No text matches found."


async def extract_facts(doc_id: str, question: str | None = None) -> str:
    """Extract key facts from document."""
    chunks = await db.get_doc_chunks(doc_id)
    # Extract from chunk metadata
    facts = []
    for c in chunks:
        content = c.get("content", "")
        if question and question.lower() not in content.lower():
            continue
        if "fact" in c.get("chunk_type", "").lower():
            facts.append(content[:200])
        elif len(content.strip()) > 50:
            facts.append(content[:200])
    
    seen = set()
    unique = []
    for f in facts:
        key = f[:40]
        if key not in seen:
            seen.add(key)
            unique.append(f)
    
    return "\n".join(f"I: {f}" for f in unique[:20])


async def get_polygon_view(doc_id: str, page_num: int, block_index: int | None = None) -> dict:
    """Get page with polygon coordinates for visualization."""
    chunks = await db.get_doc_chunks(doc_id)
    page_chunks = [c for c in chunks if c.get("page_num") == page_num]
    
    blocks = []
    for c in page_chunks:
        poly = c.get("polygon")
        if poly and isinstance(poly, dict):
            blocks.append({
                "index": c.get("block_index", 0),
                "page": page_num,
                "polygon": poly,
                "content": c.get("content", "")[:200],
                "type": c.get("chunk_type", "text"),
            })
        elif block_index is None or c.get("block_index") == block_index:
            blocks.append({
                "index": c.get("block_index", 0),
                "page": page_num,
                "polygon": {"bbox": [0, 0, 1000, 1000]},
                "content": c.get("content", "")[:200],
                "type": c.get("chunk_type", "text"),
            })
    
    # Compute page bounding box from all polygons
    all_polys = [b["polygon"] for b in blocks if b["polygon"] and isinstance(b["polygon"], dict)]
    
    pages_data = {"page_num": page_num, "blocks": blocks}
    
    if all_polys:
        min_x = min(p.get("bbox", [0, 0, 0, 0])[:2] for p in all_polys)
        max_x = max(p.get("bbox", [0, 0, 0, 0])[2:] for p in all_polys)
        pages_data["page_bbox"] = {
            "x_min": min(min_x),
            "y_min": min(min_x, key=lambda x: x[1])[1],
            "width": max(max_x) - min(min_x),
            "height": max(max_x, key=lambda x: x[1])[1] - min(min_x, key=lambda x: x[1])[1],
        }
    
    return pages_data


async def get_memories(session_id: str) -> list[dict]:
    """Get research memories from pgvector."""
    return await db.get_memories(session_id)


async def memory_similarity_search(query: str, limit: int = 10) -> str:
    """Search across all research memories."""
    results = await db.memory_similarity_search(query, limit=limit)
    parts = []
    for r in results:
        parts.append(f"[{r.get('importance', '?')}] ({r.get('memory_type', '?')}) {r['content'][:200]}")
    return "\n".join(parts) if parts else "No memories found."


async def list_documents(status: str | None = None) -> str:
    """List all indexed documents."""
    docs = await db.list_documents(status=status)
    parts = []
    for d in docs:
        parts.append(f"- [{d['status']}] {d['filename']} (p{d.get('page_count', '?')}, {d['created_at']})")
    return "\n".join(parts) if parts else "No documents found."


async def save_finding(session_id: str, question: str, answer: str, 
                       summary: str, confidence: float = 0.8, 
                       relevant_chunks: list | None = None) -> str:
    """Save a research finding to the database."""
    finding_id = await db.store_finding(
        session_id, question, answer, summary,
        "researcher", confidence, relevant_chunks
    )
    await db.update_session(session_id, findings={"count": 1})
    return finding_id


async def save_structured_evidence(
    session_id: str, rows: list[dict]
) -> list[str]:
    """Save evidence-extraction rows with traceability fields."""
    ids = await db.save_structured_evidence(session_id, rows)
    return ids


async def get_pipeline_state(session_id: str) -> list[dict]:
    """Get saved pipeline intermediate stages."""
    return await db.get_pipeline_stages(session_id)


async def merge_evidence(rows_a: list[dict], rows_b: list[dict]) -> list[dict]:
    """Merge two evidence sets, flagging agreement/conflict by topic."""
    by_topic: dict[str, list] = {}
    for r in rows_a + rows_b:
        topic = r.get("topic", "unknown")
        by_topic.setdefault(topic, []).append(r)
    merged = []
    for topic, rs in by_topic.items():
        if len(rs) == 2:
            merged.append({
                "topic": topic,
                "source_a": rs[0].get("evidence", ""),
                "source_b": rs[1].get("evidence", ""),
                "agreement": "Yes" if rs[0].get("evidence") == rs[1].get("evidence") else "No",
                "conflict": "Yes" if rs[0].get("confidence") != rs[1].get("confidence") else "No",
                "notes": f"Combined from {rs[0].get('source_doc','')} and {rs[1].get('source_doc','')}",
            })
        else:
            merged.append({k: rs[0].get(k, "") for k in ("topic", "description", "trace_ref", "evidence", "confidence")})
    return merged


# Tool registry
TOOLS = {
    "read_document": read_document,
    "read_chunks": read_chunks,
    "get_page_text": get_page_text,
    "vector_search": vector_search,
    "text_search": text_search,
    "extract_facts": extract_facts,
    "get_polygon_view": get_polygon_view,
    "get_memories": get_memories,
    "memory_similarity_search": memory_similarity_search,
    "list_documents": list_documents,
    "save_finding": save_finding,
    "save_structured_evidence": save_structured_evidence,
    "get_pipeline_state": get_pipeline_state,
    "merge_evidence": merge_evidence,
}