AI/crawler_platform/app/api/routes.py

from __future__ import annotations

from dataclasses import asdict
from typing import Any

from fastapi import HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import select

from crawler_platform.app.config.loader import load_project_config
from crawler_platform.app.core.crawler.discovery import discover_links
from crawler_platform.app.core.crawler.fetchers import RobotsPolicy, make_fetcher
from crawler_platform.app.core.crawler.pipeline import CrawlPipeline
from crawler_platform.app.core.crawler.site_crawler import SiteCrawler
from crawler_platform.app.core.database import models
from crawler_platform.app.core.database.repository import KnowledgeRepository
from crawler_platform.app.core.database.session import session_scope
from crawler_platform.app.core.extractor.ai_provider import list_openai_compatible_models
from crawler_platform.app.core.extractor.factory import extractor_for_domain
from crawler_platform.app.core.ontology.definitions import ontology_for_domain
from crawler_platform.app.core.ontology.mapper import ontology_to_dict
from crawler_platform.app.core.recommendation.scorer import PreferenceInput, RuleBasedRecommender


class CrawlRequest(BaseModel):
    config_path: str
    source_name: str
    url: str
    extractor_provider: str = "rule_based"
    extractor_model: str | None = None
    extractor_base_url: str | None = None


class SiteCrawlRequest(CrawlRequest):
    max_depth: int = 2
    max_pages: int = 50
    same_domain_only: bool = True
    analyze_page_types: list[str] = Field(default_factory=lambda: ["product", "brand", "review"])


class DiscoverRequest(BaseModel):
    config_path: str
    source_name: str
    url: str
    limit: int = 30


class RecommendRequest(BaseModel):
    project_name: str
    target_entity_type: str = "Perfume"
    preferences: dict = Field(default_factory=dict)
    limit: int = 10


class CreateProjectRequest(BaseModel):
    config_path: str


class UpdateClaimConfidenceRequest(BaseModel):
    confidence: float
    reason: str | None = None


class MergeEntitiesRequest(BaseModel):
    project_name: str
    source_entity_id: int
    target_entity_id: int


class ExtractorModelsRequest(BaseModel):
    provider: str
    base_url: str | None = None


def register_routes(app, database_url: str) -> None:
    @app.get("/health")
    def health():
        return {"ok": True}

    @app.get("/projects")
    def projects():
        with session_scope(database_url) as session:
            rows = session.scalars(select(models.Project).order_by(models.Project.created_at.desc())).all()
            return [
                {
                    "id": project.id,
                    "name": project.name,
                    "domain": project.domain,
                    "created_at": project.created_at.isoformat(),
                    "updated_at": project.updated_at.isoformat(),
                }
                for project in rows
            ]

    @app.post("/projects")
    def create_project(request: CreateProjectRequest):
        config = load_project_config(request.config_path)
        with session_scope(database_url) as session:
            project = KnowledgeRepository(session).upsert_project(config)
            return {"id": project.id, "name": project.name, "domain": project.domain}

    @app.get("/projects/{project_name}")
    def project_detail(project_name: str):
        with session_scope(database_url) as session:
            project = KnowledgeRepository(session).get_project(project_name)
            sources = session.scalars(select(models.Source).where(models.Source.project_id == project.id)).all()
            return {
                "id": project.id,
                "name": project.name,
                "domain": project.domain,
                "config": project.config,
                "sources": [
                    {
                        "id": source.id,
                        "name": source.name,
                        "type": source.type,
                        "trust_level": source.trust_level,
                        "respect_robots_txt": source.respect_robots_txt,
                        "rate_limit_per_minute": source.rate_limit_per_minute,
                    }
                    for source in sources
                ],
            }

    @app.get("/ontology/{domain}")
    def ontology(domain: str):
        return ontology_to_dict(ontology_for_domain(domain))

    @app.post("/extractors/models")
    def extractor_models(request: ExtractorModelsRequest):
        try:
            if request.provider == "lm_studio":
                models = list_openai_compatible_models(request.base_url or "http://localhost:1234/v1")
                return {"ok": True, "models": models}
            if request.provider == "openai":
                import os

                models = list_openai_compatible_models(
                    request.base_url or "https://api.openai.com/v1",
                    os.getenv("OPENAI_API_KEY"),
                )
                return {"ok": True, "models": models}
            if request.provider == "ollama":
                return {"ok": False, "error": "Ollama model listing is not implemented yet. Enter the model manually."}
            return {"ok": True, "models": [{"id": "rule_based", "owned_by": "local"}]}
        except Exception as exc:
            return {"ok": False, "error": str(exc), "models": []}

    @app.post("/crawl")
    def crawl(request: CrawlRequest):
        config = load_project_config(request.config_path)
        with session_scope(database_url) as session:
            repo = KnowledgeRepository(session)
            pipeline = CrawlPipeline(
                repo,
                extractor_for_domain(
                    config.domain,
                    provider=request.extractor_provider,
                    model=request.extractor_model,
                    base_url=request.extractor_base_url,
                ),
            )
            try:
                result = pipeline.crawl_url(config, request.source_name, request.url)
            except Exception as exc:
                raise HTTPException(status_code=400, detail=str(exc)) from exc
            return {"page_id": result.page_id, "claim_count": result.claim_count, "entity_count": result.entity_count}

    @app.post("/crawl-site")
    def crawl_site(request: SiteCrawlRequest):
        config = load_project_config(request.config_path)
        with session_scope(database_url) as session:
            repo = KnowledgeRepository(session)
            crawler = SiteCrawler(
                repo,
                extractor_for_domain(
                    config.domain,
                    provider=request.extractor_provider,
                    model=request.extractor_model,
                    base_url=request.extractor_base_url,
                ),
            )
            try:
                result = crawler.crawl_site(
                    config,
                    request.source_name,
                    request.url,
                    max_depth=max(request.max_depth, 0),
                    max_pages=max(min(request.max_pages, 500), 1),
                    same_domain_only=request.same_domain_only,
                    analyze_page_types=set(request.analyze_page_types),
                )
            except Exception as exc:
                raise HTTPException(status_code=400, detail=str(exc)) from exc
            return asdict(result)

    @app.post("/discover")
    def discover(request: DiscoverRequest):
        config = load_project_config(request.config_path)
        source_config = config.source_by_name(request.source_name)
        robots = RobotsPolicy()
        if not robots.allowed(request.url, source_config.respect_robots_txt):
            return {"ok": False, "error": "robots.txt does not allow discovery for this URL", "links": []}
        fetcher = make_fetcher(source_config.fetcher, source_config.rate_limit_per_minute)
        result = fetcher.fetch(request.url)
        links = discover_links(result.html, result.final_url or request.url, request.limit)
        return {
            "ok": True,
            "status_code": result.status_code,
            "final_url": result.final_url,
            "links": [asdict(link) for link in links],
        }

    @app.get("/projects/{project_name}/entities")
    def project_entities(project_name: str, entity_type: str | None = None, limit: int = 50):
        with session_scope(database_url) as session:
            project = KnowledgeRepository(session).get_project(project_name)
            query = select(models.Entity).where(models.Entity.project_id == project.id)
            if entity_type:
                query = query.where(models.Entity.entity_type == entity_type)
            entities = session.scalars(query.limit(limit)).all()
            return [
                {
                    "id": entity.id,
                    "type": entity.entity_type,
                    "name": entity.name,
                    "metadata": entity.metadata_json,
                }
                for entity in entities
            ]

    @app.get("/projects/{project_name}/claims")
    def project_claims(project_name: str, limit: int = 100):
        with session_scope(database_url) as session:
            project = KnowledgeRepository(session).get_project(project_name)
            rows = session.execute(
                select(models.Claim, models.Source, models.Page, models.Entity)
                .join(models.Source, models.Claim.source_id == models.Source.id)
                .join(models.Page, models.Claim.page_id == models.Page.id, isouter=True)
                .join(models.Entity, models.Claim.subject_entity_id == models.Entity.id)
                .where(models.Claim.project_id == project.id)
                .order_by(models.Claim.last_seen_at.desc())
                .limit(limit)
            ).all()
            results: list[dict[str, Any]] = []
            for claim, source, page, subject in rows:
                object_name = None
                if claim.object_entity_id:
                    object_entity = session.get(models.Entity, claim.object_entity_id)
                    object_name = object_entity.name if object_entity else None
                evidence = session.scalar(
                    select(models.Evidence)
                    .where(models.Evidence.claim_id == claim.id)
                    .order_by(models.Evidence.created_at.desc())
                )
                results.append(
                    {
                        "id": claim.id,
                        "subject": subject.name,
                        "subject_type": subject.entity_type,
                        "predicate": claim.predicate,
                        "object": object_name,
                        "object_value": claim.object_value,
                        "source": source.name,
                        "page_url": page.url if page else None,
                        "confidence": claim.confidence,
                        "confidence_reason": claim.confidence_reason,
                        "evidence_text": evidence.evidence_text if evidence else None,
                        "last_seen_at": claim.last_seen_at.isoformat(),
                    }
                )
            return results

    @app.patch("/claims/{claim_id}/confidence")
    def update_claim_confidence(claim_id: int, request: UpdateClaimConfidenceRequest):
        confidence = min(max(request.confidence, 0.0), 1.0)
        with session_scope(database_url) as session:
            claim = session.get(models.Claim, claim_id)
            if claim is None:
                return {"ok": False, "error": "claim not found"}
            claim.confidence = confidence
            claim.confidence_reason = request.reason or "manual admin update"
            claim.last_seen_at = models.utcnow()
            return {"ok": True, "claim_id": claim.id, "confidence": claim.confidence}

    @app.post("/entities/merge")
    def merge_entities(request: MergeEntitiesRequest):
        with session_scope(database_url) as session:
            project = KnowledgeRepository(session).get_project(request.project_name)
            source = session.get(models.Entity, request.source_entity_id)
            target = session.get(models.Entity, request.target_entity_id)
            if source is None or target is None or source.project_id != project.id or target.project_id != project.id:
                return {"ok": False, "error": "entity not found in project"}
            session.query(models.Claim).filter(models.Claim.subject_entity_id == source.id).update(
                {models.Claim.subject_entity_id: target.id}
            )
            session.query(models.Claim).filter(models.Claim.object_entity_id == source.id).update(
                {models.Claim.object_entity_id: target.id}
            )
            session.query(models.Relation).filter(models.Relation.subject_entity_id == source.id).update(
                {models.Relation.subject_entity_id: target.id}
            )
            session.query(models.Relation).filter(models.Relation.object_entity_id == source.id).update(
                {models.Relation.object_entity_id: target.id}
            )
            source.metadata_json = {**(source.metadata_json or {}), "merged_into": target.id}
            source.updated_at = models.utcnow()
            return {"ok": True, "source_entity_id": source.id, "target_entity_id": target.id}

    @app.get("/projects/{project_name}/recommendation-tags")
    def recommendation_tags(project_name: str):
        with session_scope(database_url) as session:
            project = KnowledgeRepository(session).get_project(project_name)
            tag_predicates = {
                "hasTopNote",
                "hasMiddleNote",
                "hasBaseNote",
                "hasScentNote",
                "hasFlavorNote",
                "evokesMood",
                "suitableForSeason",
                "suitableForOccasion",
                "hasReviewKeyword",
            }
            rows = session.execute(
                select(models.Claim, models.Entity)
                .join(models.Entity, models.Claim.object_entity_id == models.Entity.id)
                .where(models.Claim.project_id == project.id, models.Claim.predicate.in_(tag_predicates))
            ).all()
            grouped: dict[str, dict[str, Any]] = {}
            for claim, entity in rows:
                key = f"{claim.predicate}:{entity.canonical_name}"
                if key not in grouped:
                    grouped[key] = {
                        "predicate": claim.predicate,
                        "name": entity.name,
                        "type": entity.entity_type,
                        "support_count": 0,
                        "max_confidence": 0.0,
                    }
                grouped[key]["support_count"] += 1
                grouped[key]["max_confidence"] = max(grouped[key]["max_confidence"], claim.confidence)
            return sorted(grouped.values(), key=lambda item: (item["predicate"], -item["support_count"], item["name"]))

    @app.post("/recommend")
    def recommend(request: RecommendRequest):
        with session_scope(database_url) as session:
            repo = KnowledgeRepository(session)
            project = repo.get_project(request.project_name)
            recommender = RuleBasedRecommender(session)
            pref = PreferenceInput(**request.preferences)
            items = recommender.recommend(project.id, request.target_entity_type, pref, request.limit)
            return [asdict(item) for item in items]