# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import json
import os
import random
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, Optional

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from uvicorn.config import LOGGING_CONFIG

# Локальные импорты из вашего оригинального файла
import items
from config import get_config
from frame.clients import Client, HuggingFaceClient, OpenAIClient
from frame.harness4 import FrameConfigV4, FrameV4
from frame.trace import Trace
from scan_research import do_reporting as real_reporting
from scan_research import do_research as real_research
from scan_research import generate_session_key
from scan_research_dry import do_reporting as dry_reporting
from scan_research_dry import do_research as dry_research

# Получаем конфигурацию
config = get_config()

# ========================================================================================
#                    ИЗМЕНЕНИЕ ДЛЯ HUGGING FACE SPACES
#
# 1. Создаем отдельное приложение (sub-app) для API.
# ========================================================================================
api_app = FastAPI(
    title="Universal Deep Research Backend API",
    description="Intelligent research and reporting service using LLMs and web search",
    version="1.0.0",
)

# Настройка логирования
LOGGING_CONFIG["formatters"]["default"]["fmt"] = "%(asctime)s [%(name)s] %(levelprefix)s %(message)s"

# Настройка CORS (можно удалить, если развертывание в одном домене)
api_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Разрешаем все источники для простоты
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Модели Pydantic из вашего файла
class Message(BaseModel):
    text: str

class ResearchRequest(BaseModel):
    dry: bool = False
    session_key: Optional[str] = None
    start_from: str = "research"
    strategy_id: Optional[str] = None
    strategy_content: Optional[str] = None
    prompt: Optional[str] = None
    mock_directory: str = "mock_instances/stocks_24th_3_sections"

# Вспомогательные функции из вашего файла
def build_events_path(session_key: str) -> str:
    return f"instances/{session_key}.events.jsonl"

def make_message(
    event: Dict[str, Any],
    session_key: str | None = None,
    timestamp_the_event: bool = True,
) -> str:
    if timestamp_the_event:
        event = {**event, "timestamp": datetime.now().isoformat()}
    if session_key:
        items.register_item(build_events_path(session_key), event)
    return json.dumps({"event": event, "session_key": session_key}) + "\n"

# ========================================================================================
#                    ИЗМЕНЕНИЕ ДЛЯ HUGGING FACE SPACES
#
# 2. Все эндпоинты @app.post(...) заменяем на @api_app.post(...)
# ========================================================================================

@api_app.post("/research")
async def start_research(request: ResearchRequest):
    if request.start_from not in ["research", "reporting"]:
        raise HTTPException(status_code=400, detail="start_from must be either 'research' or 'reporting'")
    if request.start_from == "reporting" and not request.session_key:
        raise HTTPException(status_code=400, detail="session_key is required when starting from reporting phase")
    if request.start_from == "research" and not request.prompt:
        raise HTTPException(status_code=400, detail="prompt is required when starting from research phase")
    
    mock_dir = request.mock_directory or config.research.mock_directory
    research_impl = (lambda session_key, prompt: dry_research(session_key, prompt, mock_dir)) if request.dry else real_research
    reporting_impl = (lambda session_key: dry_reporting(session_key, mock_dir)) if request.dry else real_reporting
    session_key = request.session_key or generate_session_key()
    research_gen = research_impl(session_key, request.prompt) if request.start_from == "research" else None
    reporting_gen = reporting_impl(session_key)
    
    return StreamingResponse(
        stream_research_events(research_gen, reporting_gen, request.start_from == "research", session_key),
        media_type="application/x-ndjson",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "Content-Encoding": "none"},
    )

@api_app.post("/research2")
async def start_research2(request: ResearchRequest):
    if request.start_from not in ["research"]:
        raise HTTPException(status_code=400, detail="start_from must be 'research'")
    if request.start_from == "research" and not request.prompt:
        raise HTTPException(status_code=400, detail="prompt is required when starting from research phase")

    session_key = generate_session_key()
    if request.strategy_id is None or request.strategy_id == "default":
        research_impl = (lambda session_key, prompt: dry_research(session_key, prompt, "mock_instances/stocks_24th_3_sections")) if request.dry else real_research
        reporting_impl = (lambda session_key: dry_reporting(session_key, "mock_instances/stocks_24th_3_sections")) if request.dry else real_reporting
        session_key = request.session_key or generate_session_key()
        research_gen = research_impl(session_key, request.prompt) if request.start_from == "research" else None
        reporting_gen = reporting_impl(session_key)
        return StreamingResponse(
            stream_research_events(research_gen, reporting_gen, request.start_from == "research", session_key),
            media_type="application/x-ndjson",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "Content-Encoding": "none"},
        )
    
    return StreamingResponse(
        stream_research2_events(session_key, request.prompt, request.strategy_id, request.strategy_content),
        media_type="application/x-ndjson",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "Content-Encoding": "none"},
    )

# Асинхронные генераторы событий остаются без изменений
async def stream_research_events(
    research_fn: AsyncGenerator[Dict[str, Any], None],
    reporting_fn: AsyncGenerator[Dict[str, Any], None],
    do_research: bool,
    session_key: str,
) -> AsyncGenerator[str, None]:
    try:
        yield make_message({"type": "started", "description": "Waking up the Deep Research Backend"}, session_key)
        error_event_encountered = False
        if do_research:
            async for event in research_fn:
                if event["type"] == "error":
                    error_event_encountered = True
                yield make_message(event, session_key)
        if not error_event_encountered:
            async for event in reporting_fn:
                yield make_message(event, session_key)
            yield make_message({"type": "completed", "description": "Research and reporting completed"}, session_key)
    except asyncio.CancelledError:
        yield make_message({"type": "cancelled", "description": "Research was cancelled"}, session_key)
        raise

async def stream_research2_events(
    session_key: str, prompt: str, strategy_id: str, strategy_content: str
) -> AsyncGenerator[str, None]:
    try:
        yield make_message({"type": "started", "description": "Waking up the Universal Deep Research Backend"}, session_key)
        random.seed(config.research.random_seed)
        comm_trace_timestamp = datetime.now().strftime("%Y%m%d_%H-%M-%S")
        comm_trace_filename = f"{config.logging.log_dir}/comms_{comm_trace_timestamp}.log"
        comm_trace = Trace(comm_trace_filename, copy_into_stdout=config.logging.copy_into_stdout)
        client: Client = OpenAIClient(base_url="https://integrate.api.nvidia.com/v1", model="nvdev/meta/llama-3.1-70b-instruct", trace=comm_trace)
        frame_config = FrameConfigV4(
            long_context_cutoff=config.frame.long_context_cutoff,
            force_long_context=config.frame.force_long_context,
            max_iterations=config.frame.max_iterations,
            interaction_level=config.frame.interaction_level,
        )
        harness = FrameV4(client_profile=client, errand_profile={}, compilation_trace=True, execution_trace="file_and_stdout")
        messages = []
        preamble_files = ["frame/prompts/udr_minimal_generating/0.code_skill.py"]
        for path in preamble_files:
            type = path.split(".")[-2]
            with open(path, "r") as f:
                messages.append({"mid": len(messages), "role": "user", "content": f.read(), "type": type})
        messages.append({"mid": len(messages), "role": "user", "content": "The following is the prompt data to be used in later procedures.\n\nPROMPT:\n" + prompt, "type": "data"})
        messages.append({"mid": len(messages), "role": "user", "content": strategy_content, "type": "generating_routine"})
        for i in range(len(messages)):
            messages_so_far = messages[: i + 1]
            yield make_message({"type": "generic", "description": f"Processing agentic instructions: {i + 1} of {len(messages)}"}, session_key)
            for notification in harness.generate_with_notifications(messages=messages_so_far, frame_config=frame_config):
                yield make_message(notification, session_key)
        yield make_message({"type": "completed", "description": "Research completed"}, session_key)
    except asyncio.CancelledError:
        yield make_message({"type": "cancelled", "description": "Research was cancelled"}, session_key)
        raise

# ========================================================================================
#                    ИЗМЕНЕНИЕ ДЛЯ HUGGING FACE SPACES
#
# 3. Создаем главное приложение `app`.
# 4. Монтируем `api_app` на `/api`.
# 5. Монтируем статический фронтенд в корень `/`.
# ========================================================================================

app = FastAPI()

# Монтируем API
app.mount("/api", api_app)

# Монтируем статический фронтенд
# Это должно быть в самом конце файла!
app.mount("/", StaticFiles(directory="/app/static_frontend", html=True), name="static")