import gradio as gr
import threading
import json
import asyncio
import sys
import io
import threading as _threading
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel as _BaseModel
from agent.crew import run_research
from output.exporter import export_markdown, export_pdf

_current_result: dict = {}
_lock = threading.Lock()

DEPTH_DESCRIPTIONS = {
    "quick": "Quick (5 sources, ~30s) — surface overview",
    "standard": "Standard (10 sources, ~90s) — balanced depth",
    "deep": "Deep (15 sources, ~3min) — exhaustive, academic-heavy",
}

CSS = """
#report-box { font-family: 'Georgia', serif; font-size: 14px; }
#confidence-badge { font-weight: bold; font-size: 16px; padding: 8px; border-radius: 6px; }
"""


def _confidence_color(report: str) -> str:
    import re
    match = re.search(r"Confidence Assessment.*?(HIGH|MEDIUM|LOW|INSUFFICIENT)", report, re.DOTALL | re.IGNORECASE)
    label = match.group(1).upper() if match else "UNKNOWN"
    colors = {"HIGH": "#d4edda", "MEDIUM": "#fff3cd", "LOW": "#f8d7da", "INSUFFICIENT": "#e2e3e5", "UNKNOWN": "#e2e3e5"}
    return f'<div style="background:{colors[label]};padding:10px;border-radius:6px;font-weight:bold">Confidence: {label}</div>'


def run_research_ui(query: str, depth_label: str, progress=gr.Progress()):
    global _current_result

    depth = [k for k, v in DEPTH_DESCRIPTIONS.items() if v == depth_label]
    depth = depth[0] if depth else "standard"

    if not query.strip():
        return "Please enter a research topic.", "", None, None

    progress(0.05, desc="Starting research crew...")

    try:
        progress(0.15, desc="Searching sources...")
        result = run_research(query=query.strip(), depth=depth)
        progress(0.90, desc="Formatting report...")

        with _lock:
            _current_result = result

        report = result.get("report", "No report generated.")
        badge = _confidence_color(report)

        sources = result.get("sources", [])
        source_summary = ""
        if sources:
            source_summary = f"**{len(sources)} accepted sources** (sorted by credibility)\n\n"
            for s in sources[:8]:
                score = s.get("credibility_score", 0)
                conf = s.get("confidence", "low")
                title = s.get("title", "Untitled")[:60]
                url = s.get("url", "")
                source_summary += f"- [{title}]({url}) — score: **{score:.2f}** ({conf})\n"

        progress(1.0, desc="Done.")
        return report, badge + "\n\n" + source_summary, None, None

    except Exception as e:
        return f"Error during research: {e}", "", None, None


def export_md():
    with _lock:
        result = dict(_current_result)
    if not result.get("report"):
        return None
    return export_markdown(result["report"], result.get("query", "research"))


def export_pdf_fn():
    with _lock:
        result = dict(_current_result)
    if not result.get("report"):
        return None
    return export_pdf(result["report"], result.get("query", "research"))


with gr.Blocks(css=CSS, title="AI Research Agent") as demo:
    gr.Markdown(
        """
# AI Research Agent
**Powered by CrewAI · Tavily · arXiv · Groq**

Searches multiple authoritative sources, scores credibility, and synthesizes a fully cited report.
Only sources scoring ≥ 0.35/1.0 on the credibility scale are included.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            query_input = gr.Textbox(
                label="Research Topic",
                placeholder="e.g. What are the latest advances in mRNA vaccine technology?",
                lines=2,
            )
        with gr.Column(scale=1):
            depth_input = gr.Radio(
                choices=list(DEPTH_DESCRIPTIONS.values()),
                value=DEPTH_DESCRIPTIONS["standard"],
                label="Research Depth",
            )

    run_btn = gr.Button("Start Research", variant="primary", size="lg")

    with gr.Tabs():
        with gr.TabItem("Report"):
            report_output = gr.Markdown(label="Research Report", elem_id="report-box")
        with gr.TabItem("Sources & Confidence"):
            meta_output = gr.Markdown(label="Source Summary")

    with gr.Row():
        md_btn = gr.Button("Export Markdown")
        pdf_btn = gr.Button("Export PDF")

    md_file = gr.File(label="Markdown Download", visible=True)
    pdf_file = gr.File(label="PDF Download", visible=True)

    run_btn.click(
        fn=run_research_ui,
        inputs=[query_input, depth_input],
        outputs=[report_output, meta_output, md_file, pdf_file],
    )

    md_btn.click(fn=export_md, inputs=[], outputs=[md_file])
    pdf_btn.click(fn=export_pdf_fn, inputs=[], outputs=[pdf_file])

    gr.Examples(
        examples=[
            ["What are the health risks of microplastics in drinking water?", DEPTH_DESCRIPTIONS["standard"]],
            ["Latest advances in large language model reasoning capabilities 2024-2025", DEPTH_DESCRIPTIONS["deep"]],
            ["Economic impact of remote work on urban commercial real estate", DEPTH_DESCRIPTIONS["standard"]],
            ["CRISPR gene editing applications in treating genetic diseases", DEPTH_DESCRIPTIONS["quick"]],
        ],
        inputs=[query_input, depth_input],
    )


# ── Attach REST API to Gradio's underlying FastAPI app ────────────

demo.app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

_thread_local_app = _threading.local()
_original_stdout = sys.stdout


class _RoutedStream(io.TextIOBase):
    def write(self, text: str) -> int:
        q = getattr(_thread_local_app, "log_queue", None)
        if q is not None:
            for line in text.splitlines():
                line = line.strip()
                if line and not line.startswith("\x1b"):
                    try:
                        q.put_nowait(line)
                    except Exception:
                        pass
        else:
            try:
                _original_stdout.write(text)
            except Exception:
                pass
        return len(text)

    def flush(self):
        try:
            _original_stdout.flush()
        except Exception:
            pass

    def isatty(self):
        return False


sys.stdout = _RoutedStream(_original_stdout)


class _ResearchReq(_BaseModel):
    query: str
    depth: str = "standard"


async def _sse_stream(query: str, depth: str):
    loop = asyncio.get_event_loop()
    log_queue: asyncio.Queue = asyncio.Queue(maxsize=500)
    done_event = asyncio.Event()
    result_holder: dict = {}

    def _worker():
        _thread_local_app.log_queue = log_queue
        try:
            result = run_research(query, depth)
            result_holder.update(result)
        except Exception as e:
            result_holder["error"] = str(e)
        finally:
            _thread_local_app.log_queue = None
            loop.call_soon_threadsafe(done_event.set)

    _threading.Thread(target=_worker, daemon=True).start()

    while not done_event.is_set():
        try:
            line = log_queue.get_nowait()
            yield f"data: {json.dumps(str(line))}\n\n"
        except Exception:
            await asyncio.sleep(0.05)

    while not log_queue.empty():
        try:
            line = log_queue.get_nowait()
            yield f"data: {json.dumps(str(line))}\n\n"
        except Exception:
            break

    if "error" in result_holder:
        yield f"data: {json.dumps('[ERROR] ' + result_holder['error'])}\n\n"
    else:
        yield f"data: {json.dumps({'type': 'report', 'content': result_holder.get('report', '')})}\n\n"
        yield f"data: {json.dumps({'type': 'sources', 'content': result_holder.get('sources', [])})}\n\n"

    yield 'data: "__DONE__"\n\n'


@demo.app.post("/research")
async def research_api(req: _ResearchReq):
    if not req.query.strip():
        return {"error": "Query cannot be empty"}
    return StreamingResponse(
        _sse_stream(req.query, req.depth),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


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


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)