#!/usr/bin/env python3
"""
AI Research Engine — Backend API
Runs on CT 145, handles all heavy lifting.
Exposes REST API consumed by the MCP proxy (MacBook) and dashboard.
"""

import os
import json
import httpx
from fastapi import FastAPI, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
import uvicorn

app = FastAPI(title="AI Research Engine Backend")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# ── Config ──────────────────────────────────────────────────
YACY_URL       = os.getenv("YACY_URL", "http://localhost:8090")
OPENSEARCH_URL = os.getenv("OPENSEARCH_URL", "http://localhost:9200")
QDRANT_URL     = os.getenv("QDRANT_URL", "http://10.30.20.68:6333")
OLLAMA_URL     = os.getenv("OLLAMA_URL", "http://10.30.20.186:11434")
OLLAMA_MODEL   = os.getenv("OLLAMA_MODEL", "ornith:latest")
INDEX_NAME     = os.getenv("INDEX_NAME", "research_docs")

client = httpx.Client(timeout=30.0)
ollama = httpx.Client(timeout=120.0, base_url=OLLAMA_URL)

# ── Helpers ──────────────────────────────────────────────────

def _ensure_index():
    try:
        client.get(f"{OPENSEARCH_URL}/{INDEX_NAME}")
    except Exception:
        try:
            client.put(f"{OPENSEARCH_URL}/{INDEX_NAME}", json={
                "settings": {"number_of_shards": 1, "number_of_replicas": 0},
                "mappings": {"properties": {
                    "url": {"type": "keyword"}, "title": {"type": "text"},
                    "content": {"type": "text"}, "excerpt": {"type": "text"},
                    "category": {"type": "keyword"}, "source_domain": {"type": "keyword"},
                    "crawled_at": {"type": "date"}, "indexed_at": {"type": "date"},
                    "metadata": {"type": "object"},
                }}
            })
        except Exception:
            pass

def _ensure_qdrant():
    try:
        client.get(f"{QDRANT_URL}/collections/{INDEX_NAME}")
    except Exception:
        try:
            client.put(f"{QDRANT_URL}/collections/{INDEX_NAME}", json={
                "vectors": {"size": 768, "distance": "Cosine"}
            })
        except Exception:
            pass

def _ai_chat(prompt: str, system: str = "You are a research assistant. Be concise and factual.") -> str:
    r = ollama.post("/api/chat", json={
        "model": OLLAMA_MODEL, "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": prompt},
        ], "stream": False,
        "options": {"temperature": 0.3, "num_predict": 2048},
    })
    body = r.json()
    return body.get("message", {}).get("content", "") or body.get("thinking", "") or ""

def _get_embedding(text: str) -> list:
    try:
        r = ollama.post("/api/embeddings", json={"model": "nomic-embed-text-v2-moe:latest", "prompt": text[:2048]})
        return r.json().get("embedding", [])
    except Exception:
        return []


# ── Status ────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok"}

@app.get("/api/status")
def status():
    svc = {}
    try:
        r = client.get(f"{OPENSEARCH_URL}/_cluster/health")
        cnt = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_count").json() if r.status_code == 200 else {"count": 0}
        svc["opensearch"] = {"status": r.json().get("status", "down"), "documents": cnt.get("count", 0)}
    except Exception:
        svc["opensearch"] = {"status": "down"}
    try:
        r = client.get(f"{QDRANT_URL}/healthz")
        col = client.get(f"{QDRANT_URL}/collections/{INDEX_NAME}").json()
        svc["qdrant"] = {"status": "ok", "points": col.get("result", {}).get("points_count", 0)}
    except Exception:
        svc["qdrant"] = {"status": "down"}
    try:
        r = client.get(f"{YACY_URL}/yacysearch.json", params={"query": "test", "maximumRecords": 1})
        svc["yacy"] = {"status": "ok" if r.status_code == 200 else "down"}
    except Exception:
        svc["yacy"] = {"status": "down"}
    try:
        r = ollama.get("/api/tags")
        svc["ollama"] = {"status": "ok", "models": len(r.json().get("models", []))}
    except Exception:
        svc["ollama"] = {"status": "down"}
    return {"services": svc}


# ── Search ────────────────────────────────────────────────────

@app.get("/api/search")
def search_web(q: str = Query(...), category: str = "", limit: int = 10):
    _ensure_index()
    body = {
        "size": limit,
        "query": {"bool": {"must": [{"multi_match": {"query": q, "fields": ["title^3", "content", "excerpt"]}}]}},
        "highlight": {"fields": {"content": {"fragment_size": 200, "number_of_fragments": 2}}},
    }
    if category:
        body["query"]["bool"]["filter"] = [{"term": {"category": category}}]
    try:
        r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json=body, params={"refresh": "true"})
        result = r.json()
        hits = []
        for h in result.get("hits", {}).get("hits", []):
            src = h["_source"]
            hits.append({"url": src.get("url"), "title": src.get("title"),
                "excerpt": src.get("excerpt") or (h.get("highlight", {}).get("content", [""])[0]),
                "category": src.get("category"), "crawled_at": src.get("crawled_at"), "score": h["_score"]})
        return {"query": q, "total": result.get("hits", {}).get("total", {}).get("value", 0), "hits": hits}
    except Exception as e:
        return {"query": q, "total": 0, "hits": [], "note": f"Index may be empty. {e}"}


@app.get("/api/semantic-search")
def semantic_search(q: str = Query(...), limit: int = 10):
    _ensure_qdrant()
    emb = _get_embedding(q)
    if not emb:
        return {"hits": [], "error": "Embedding model not available"}
    try:
        r = client.post(f"{QDRANT_URL}/collections/{INDEX_NAME}/points/search", json={
            "vector": emb, "limit": limit, "with_payload": True, "with_vector": False})
        hits = [{"url": p.get("payload", {}).get("url"), "title": p.get("payload", {}).get("title"),
            "excerpt": str(p.get("payload", {}).get("excerpt", ""))[:300], "score": p.get("score")}
            for p in r.json().get("result", [])]
        return {"query": q, "total": len(hits), "hits": hits}
    except Exception as e:
        return {"hits": [], "error": str(e)}


# ── Crawl ─────────────────────────────────────────────────────

@app.get("/api/crawl")
def crawl_url(url: str = Query(...), depth: int = 1):
    try:
        r = client.get(f"{YACY_URL}/Crawler_p.json", params={
            "crawlingDomMaxPages": 50, "crawlingDepth": depth,
            "crawlingStart": url, "crawlingQ": "on",
            "bookmarkTitle": "research", "bookmarkFolder": "/research",
            "indexText": "on", "indexMedia": "on",
            "crawlingMode": "url", "cachePolicy": "iffresh",
        })
        return {"status": "crawl_started", "url": url, "depth": depth}
    except Exception as e:
        return {"status": "error", "url": url, "error": str(e)}


@app.get("/api/crawl-topic")
def crawl_topic(topic: str = Query(...), max_urls: int = 20):
    discovered = []
    try:
        r = client.get(f"{YACY_URL}/yacysearch.json", params={"query": topic, "maximumRecords": max_urls, "resource": "global"})
        for ch in r.json().get("channels", []):
            for item in ch.get("items", []):
                if item.get("link"):
                    discovered.append({"url": item["link"], "title": item.get("title", ""), "snippet": item.get("description", "")})
    except Exception as e:
        return {"topic": topic, "error": f"Discovery failed: {e}", "urls_discovered": 0, "urls_crawled": 0}

    crawled = 0
    for u in discovered[:max_urls]:
        try:
            client.get(f"{YACY_URL}/Crawler_p.json", params={
                "crawlingDomMaxPages": 10, "crawlingDepth": 0,
                "crawlingStart": u["url"], "crawlingQ": "on",
                "indexText": "on", "indexMedia": "on",
                "crawlingMode": "url", "cachePolicy": "iffresh",
            }, timeout=5.0)
            crawled += 1
        except Exception:
            pass

    return {"topic": topic, "urls_discovered": len(discovered), "urls_crawled": crawled}


# ── Document Retrieval ────────────────────────────────────────

@app.get("/api/document")
def retrieve_document(url: str = Query(...)):
    try:
        r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json={
            "size": 1, "query": {"term": {"url": url}}})
        hits = r.json().get("hits", {}).get("hits", [])
        if not hits:
            return {"error": "Not found", "url": url}
        src = hits[0]["_source"]
        return {"url": src.get("url"), "title": src.get("title"),
            "content": src.get("content", "")[:10000], "excerpt": src.get("excerpt"),
            "category": src.get("category"), "crawled_at": src.get("crawled_at")}
    except Exception as e:
        return {"error": str(e), "url": url}


# ── AI Synthesis ──────────────────────────────────────────────

@app.get("/api/summarize")
def summarize_sources(urls: str = Query(...), instruction: str = "Summarize key points"):
    url_list = [u.strip() for u in urls.split(",") if u.strip()]
    combined = ""
    for url in url_list[:5]:
        try:
            r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json={"size": 1, "query": {"term": {"url": url}}})
            hits = r.json().get("hits", {}).get("hits", [])
            if hits:
                src = hits[0]["_source"]
                combined += f"\n\n--- {url} ---\n{src.get('title','')}\n{src.get('content', src.get('excerpt',''))[:2000]}"
        except Exception:
            pass
    if not combined.strip():
        return {"error": "No content retrieved"}
    summary = _ai_chat(f"Instruction: {instruction}\n\nSources:{combined}\n\nProvide a structured summary.")
    return {"instruction": instruction, "sources": len(url_list), "summary": summary}


@app.get("/api/extract")
def extract_information(url: str = Query(...), schema: str = Query("company names, products, prices")):
    try:
        r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json={"size": 1, "query": {"term": {"url": url}}})
        hits = r.json().get("hits", {}).get("hits", [])
        if not hits:
            return {"error": "Not found", "url": url}
        content = hits[0]["_source"].get("content", "")[:8000]
        prompt = f"Extract: {schema}\n\nDocument:\n{content}\n\nReturn ONLY valid JSON."
        result = _ai_chat(prompt, system="Extract structured data. Return ONLY valid JSON. No explanation.")
        return {"url": url, "schema": schema, "extracted": result}
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/report")
def create_report(topic: str = Query(...), sources: str = ""):
    if sources:
        urls = [u.strip() for u in sources.split(",") if u.strip()]
    else:
        try:
            r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json={
                "size": 8, "query": {"multi_match": {"query": topic, "fields": ["title^3", "content"]}}})
            urls = [h["_source"]["url"] for h in r.json().get("hits", {}).get("hits", [])]
        except Exception:
            urls = []

    gathered = ""
    for url in urls[:8]:
        try:
            r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json={"size": 1, "query": {"term": {"url": url}}})
            hits = r.json().get("hits", {}).get("hits", [])
            if hits:
                src = hits[0]["_source"]
                gathered += f"\n\n### {src.get('title','Source')}\nURL: {url}\n{src.get('content',src.get('excerpt',''))[:1500]}"
        except Exception:
            pass

    prompt = f"""Research topic: {topic}\nSources:{gathered if gathered else ' No sources found.'}

Generate a comprehensive report:
1. Executive Summary
2. Key Findings (numbered)
3. Source Analysis
4. Knowledge Gaps
5. Recommendations

Be thorough, use markdown, cite sources."""
    report = _ai_chat(prompt, system="You are a senior research analyst. Produce thorough, structured reports.")
    return {"topic": topic, "sources_used": len(urls), "report": report}


# ── Research Pipeline ─────────────────────────────────────────

@app.get("/api/research")
def research_topic(topic: str = Query(...)):
    steps = []
    kw_result = {"hits": []}
    sem_result = {"hits": []}

    try:
        r = client.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", json={
            "size": 5, "query": {"multi_match": {"query": topic, "fields": ["title^3", "content", "excerpt"]}},
            "highlight": {"fields": {"content": {"fragment_size": 200, "number_of_fragments": 1}}},
        })
        hits = [{"title": h["_source"].get("title"), "excerpt": h.get("highlight", {}).get("content", [h["_source"].get("excerpt", "")])[0]}
                for h in r.json().get("hits", {}).get("hits", [])]
        kw_result = {"total": len(hits), "hits": hits}
        steps.append("keyword_search_done")
    except Exception:
        steps.append("keyword_search_skipped")

    try:
        emb = _get_embedding(topic)
        if emb:
            r = client.post(f"{QDRANT_URL}/collections/{INDEX_NAME}/points/search", json={
                "vector": emb, "limit": 5, "with_payload": True})
            hits2 = [{"title": p.get("payload", {}).get("title", ""), "excerpt": str(p.get("payload", {}).get("excerpt", ""))[:200],
                "score": p.get("score")} for p in r.json().get("result", [])]
            sem_result = {"total": len(hits2), "hits": hits2}
            steps.append("semantic_search_done")
    except Exception:
        steps.append("semantic_search_skipped")

    try:
        r = client.get(f"{YACY_URL}/yacysearch.json", params={"query": topic, "maximumRecords": 10, "resource": "global"})
        discovered = [item.get("link") for ch in r.json().get("channels", []) for item in ch.get("items", []) if item.get("link")]
        for url in discovered[:10]:
            try:
                client.get(f"{YACY_URL}/Crawler_p.json", params={
                    "crawlingDomMaxPages": 10, "crawlingDepth": 0, "crawlingStart": url,
                    "crawlingQ": "on", "indexText": "on", "indexMedia": "on",
                    "crawlingMode": "url", "cachePolicy": "iffresh"}, timeout=5.0)
            except Exception:
                pass
        steps.append(f"crawl_dispatched_{len(discovered[:10])}")
    except Exception:
        steps.append("crawl_skipped")

    all_sources = ""
    for h in kw_result.get("hits", [])[:3] + sem_result.get("hits", [])[:3]:
        all_sources += f"- {h.get('title', 'Unknown')}: {h.get('excerpt', '')[:200]}\n"

    summary = ""
    if all_sources:
        summary = _ai_chat(
            f"Research topic: {topic}\n\nSources:\n{all_sources}\n\nConcise research summary (3-5 paragraphs): "
            "key findings, important sources, knowledge gaps, next steps.")

    return {"topic": topic, "steps": steps,
        "keyword_results": kw_result, "semantic_results": sem_result,
        "ai_summary": summary}


# ── Dashboard ─────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
def dashboard():
    return HTMLResponse("""
<!DOCTYPE html><html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1.0">
<title>AI Research Engine</title><link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700;800&display=swap" rel="stylesheet">
<style>
:root{--bg:#09090b;--surface:#18181b;--border:#27272a;--accent:#7c3aed;--accent2:#a855f7;--text:#f4f4f5;--muted:#a1a1aa;--green:#22c55e;--red:#ef4444;--amber:#f59e0b}
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:'Inter',system-ui,sans-serif;background:var(--bg);color:var(--text);min-height:100vh}
.header{background:var(--surface);border-bottom:1px solid var(--border);padding:20px 32px;display:flex;justify-content:space-between;align-items:center;position:sticky;top:0;z-index:10}
.header h1{font-size:1.4rem;font-weight:800;background:linear-gradient(135deg,var(--accent),var(--accent2),#ec4899);-webkit-background-clip:text;-webkit-text-fill-color:transparent;letter-spacing:-0.02em}
.header .badge{font-size:0.75rem;color:var(--muted);background:var(--border);padding:6px 12px;border-radius:999px}
.container{max-width:1400px;margin:0 auto;padding:32px 24px}
.hero{text-align:center;padding:48px 0 32px}
.hero h2{font-size:2.2rem;font-weight:800;letter-spacing:-0.03em;margin-bottom:12px}
.hero p{color:var(--muted);font-size:1.1rem;max-width:600px;margin:0 auto}
.search-box{display:flex;gap:12px;max-width:800px;margin:0 auto 40px}
.search-box input{flex:1;padding:16px 24px;border-radius:16px;border:1px solid var(--border);background:var(--surface);color:var(--text);font-size:1rem;outline:none;transition:border-color .2s}
.search-box input:focus{border-color:var(--accent)}
.search-box button{padding:16px 32px;border-radius:16px;border:none;font-weight:600;font-size:1rem;cursor:pointer;transition:all .2s}
.btn-primary{background:var(--accent);color:white}
.btn-primary:hover{background:var(--accent2)}
.btn-secondary{background:var(--surface);color:var(--text);border:1px solid var(--border)}
.btn-secondary:hover{border-color:var(--accent)}
.grid{display:grid;grid-template-columns:repeat(auto-fit,minmax(260px,1fr));gap:16px;margin-bottom:32px}
.card{background:var(--surface);border:1px solid var(--border);border-radius:16px;padding:24px;transition:border-color .2s}
.card:hover{border-color:var(--accent)}
.card .label{color:var(--muted);font-size:0.8rem;text-transform:uppercase;letter-spacing:0.05em;margin-bottom:8px}
.card .value{font-size:1.8rem;font-weight:700}
.card .sub{color:var(--muted);font-size:0.85rem;margin-top:4px}
.ok{color:var(--green)}.down{color:var(--red)}.warn{color:var(--amber)}
.action-bar{display:flex;gap:10px;margin-bottom:28px;flex-wrap:wrap}
.action-bar button{padding:10px 20px;border-radius:10px;border:1px solid var(--border);background:var(--surface);color:var(--text);cursor:pointer;font-size:0.9rem;transition:all .15s}
.action-bar button:hover{border-color:var(--accent);background:#1f1f23}
.results{display:flex;flex-direction:column;gap:12px}
.result-card{background:var(--surface);border:1px solid var(--border);border-radius:14px;padding:20px;transition:border-color .15s}
.result-card:hover{border-color:var(--accent)}
.result-card h3{margin-bottom:6px;font-size:1.1rem}
.result-card h3 a{color:var(--accent2);text-decoration:none}
.result-card h3 a:hover{text-decoration:underline}
.result-card .excerpt{color:var(--muted);font-size:0.9rem;line-height:1.5}
.result-card .meta{color:#71717a;font-size:0.78rem;margin-top:10px;display:flex;gap:16px}
.loading{text-align:center;padding:48px;color:var(--muted)}
.hidden{display:none}
pre{background:var(--surface);border:1px solid var(--border);border-radius:12px;padding:20px;overflow-x:auto;font-size:0.85rem;line-height:1.6;white-space:pre-wrap}
.report{background:var(--surface);border:1px solid var(--border);border-radius:14px;padding:24px;line-height:1.7}
.report h1,.report h2,.report h3{color:var(--accent2);margin:16px 0 8px}
.report ul,.report ol{padding-left:20px;margin:8px 0}
.report li{margin:4px 0}
</style></head><body>
<div class="header"><h1>🧠 AI Research Engine</h1><div class="badge">Self-hosted · Private</div></div>
<div class="container">
<div class="hero"><h2>Your Private Research Cloud</h2><p>Discover, index, and synthesize knowledge — all on your own infrastructure.</p></div>
<div class="search-box">
<input type="text" id="query" placeholder="Research anything..." onkeydown="if(event.key==='Enter')search()">
<button class="btn-primary" onclick="search()">🔍 Search</button>
<button class="btn-secondary" onclick="research()" style="background:var(--accent);color:white;">🧪 Deep Research</button>
</div>
<div class="action-bar">
<button onclick="loadStatus()">📡 System Status</button>
<button onclick="toggleCrawl()">🕷️ Crawl URL</button>
<button onclick="loadReport()">📄 Generate Report</button>
</div>
<div id="crawl-box" class="hidden" style="margin-bottom:20px;display:flex;gap:10px;">
<input type="text" id="crawl-input" placeholder="https://example.com" style="flex:1;padding:12px 16px;border-radius:12px;border:1px solid var(--border);background:var(--surface);color:var(--text);">
<button class="btn-primary" onclick="crawl()">Crawl</button>
</div>
<div class="grid" id="stats"></div>
<div id="output"></div>
</div>
<script>
async function api(p){const r=await fetch(p);return r.json()}
async function loadStatus(){
try{const d=await api('/api/status');let h='';for(const[n,i]of Object.entries(d.services||{})){
const c=i.status==='ok'||i.status==='green'?'ok':'down';
let x='';if(i.documents!==undefined)x=i.documents+' docs';if(i.points!==undefined)x=i.points+' vectors';if(i.models!==undefined)x=i.models+' models';
h+=`<div class="card"><div class="label">${n}</div><div class="value ${c}">${i.status||'down'}</div><div class="sub">${x||'—'}</div></div>`}
document.getElementById('stats').innerHTML=h}catch(e){}
}
async function search(){
const q=document.getElementById('query').value;if(!q)return;
document.getElementById('output').innerHTML='<div class="loading">Searching...</div>';
const d=await api('/api/search?q='+encodeURIComponent(q));
let h=`<p style="color:var(--muted);margin-bottom:16px">${d.total||0} results for "${d.query||q}"</p><div class="results">`;
for(const r of(d.hits||[]))h+=`<div class="result-card"><h3><a href="${r.url||'#'}" target="_blank">${r.title||'Untitled'}</a></h3><div class="excerpt">${r.excerpt||''}</div><div class="meta"><span>⭐ ${(r.score||0).toFixed(1)}</span>${r.category?`<span>📁 ${r.category}</span>`:''}<span>${r.crawled_at||''}</span></div></div>`;
h+='</div>';document.getElementById('output').innerHTML=h}
async function research(){
const q=document.getElementById('query').value;if(!q)return;
document.getElementById('output').innerHTML='<div class="loading">Deep researching... this may take 30-60 seconds</div>';
const d=await api('/api/research?topic='+encodeURIComponent(q));
let h=`<div class="card" style="margin-bottom:16px"><div class="label">Research Complete</div><div class="value" style="font-size:1.2rem">${d.topic}</div><div class="sub">Steps: ${(d.steps||[]).join(' → ')}</div></div>`;
if(d.ai_summary)h+=`<div class="report"><h2>AI Summary</h2>${d.ai_summary.replace(/\\n/g,'<br>')}</div>`;
h+=`<div style="display:grid;grid-template-columns:1fr 1fr;gap:16px;margin-top:16px">`;
h+=`<div class="card"><div class="label">Keyword Results</div><div class="sub">${d.keyword_results?.total||0} hits</div></div>`;
h+=`<div class="card"><div class="label">Semantic Results</div><div class="sub">${d.semantic_results?.total||0} hits</div></div></div>`;
document.getElementById('output').innerHTML=h}
function toggleCrawl(){document.getElementById('crawl-box').classList.toggle('hidden')}
async function crawl(){
const u=document.getElementById('crawl-input').value;if(!u)return;
document.getElementById('output').innerHTML='<div class="loading">Dispatching crawl...</div>';
const d=await api('/api/crawl?url='+encodeURIComponent(u));
document.getElementById('output').innerHTML='<pre>'+JSON.stringify(d,null,2)+'</pre>'}
async function loadReport(){
const q=document.getElementById('query').value||'AI research';
document.getElementById('output').innerHTML='<div class="loading">Generating report...</div>';
const d=await api('/api/report?topic='+encodeURIComponent(q));
document.getElementById('output').innerHTML=`<div class="report"><h2>📊 Research Report: ${d.topic}</h2>${(d.report||'').replace(/\\n/g,'<br>')}<p style="color:var(--muted);margin-top:16px">Sources: ${d.sources_used||0}</p></div>`}
loadStatus()
</script></body></html>""")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)